Home
last modified time | relevance | path

Searched refs:mesh_split (Results 1 – 4 of 4) sorted by relevance

/aosp_15_r20/external/tensorflow/tensorflow/python/training/
H A Dslot_creator_test.py181 v = xla_sharding.mesh_split(
193 v = xla_sharding.mesh_split(
208 v = xla_sharding.mesh_split(
227 v = xla_sharding.mesh_split(
251 v = xla_sharding.mesh_split(
H A Dadam_test.py391 xla_sharding.mesh_split(
H A Dmoving_averages_test.py582 v = xla_sharding.mesh_split(v, np.array([0, 1]), [0], use_sharding_op=False)
/aosp_15_r20/external/tensorflow/tensorflow/compiler/xla/experimental/xla_sharding/
H A Dxla_sharding.py558 def mesh_split(tensor, function