Searched refs:sharding_attr (Results 1 – 3 of 3) sorted by relevance
447 sharding_attr = op.get_attr("sharding")448 grad_sharding = gen_xla_ops.xla_sharding(grad, sharding=sharding_attr)451 attr_value_pb2.AttrValue(s=sharding_attr))
849 string sharding_attr; in __anonbbc1c2d10c02() local850 TF_RETURN_IF_ERROR(c->GetAttr("manual_sharding", &sharding_attr)); in __anonbbc1c2d10c02()852 sharding.ParseFromString(sharding_attr); in __anonbbc1c2d10c02()
236 const auto& sharding_attr = sharding_attr_and_index.value(); in ExtractInputsForLogicalDevices() local242 sharding_attr.cast<mlir::StringAttr>().getValue().str()); in ExtractInputsForLogicalDevices()