Searched refs:manual_sharding (Results 1 – 3 of 3) sorted by relevance
| /external/tensorflow/tensorflow/compiler/xla/experimental/xla_sharding/ |
| D | xla_sharding.py | 370 def auto_to_manual_spmd_partition(tensor, manual_sharding): argument
386 tensor, manual_sharding=manual_sharding)
389 def manual_to_auto_spmd_partition(tensor, manual_sharding, full_shape): argument
406 tensor, manual_sharding=manual_sharding, full_shape=full_shape)
|
| /external/tensorflow/tensorflow/compiler/tf2xla/python/ |
| D | xla.py | 463 manual_sharding=op.get_attr("manual_sharding"),
471 grad, manual_sharding=op.get_attr("manual_sharding"))
|
| /external/tensorflow/tensorflow/go/op/ |
| D | wrappers.go | 46 func XlaSpmdShardToFullShape(scope *Scope, input tf.Output, manual_sharding string, full_shape tf.S…
50 attrs := map[string]interface{}{"manual_sharding": manual_sharding, "full_shape": full_shape}
27117 func XlaSpmdFullToShardShape(scope *Scope, input tf.Output, manual_sharding string) (output tf.Outp…
27121 attrs := map[string]interface{}{"manual_sharding": manual_sharding}
|