Searched refs:axis_groups (Results 1 – 6 of 6) sorted by relevance
169 SCM axis_groups = SCM_EOL; in stop_translation_timestep() local178 axis_groups = scm_cons (scm_car (s), axis_groups); in stop_translation_timestep()180 axis_groups_ = axis_groups; in stop_translation_timestep()
167 SCM axis_groups = SCM_EOL; in stop_translation_timestep() local176 axis_groups = scm_cons (scm_car (s), axis_groups); in stop_translation_timestep()178 axis_groups_ = axis_groups; in stop_translation_timestep()
548 xla.axis_groups(axis_env, tuple(out_axes.keys())))
518 replica_groups = xla.axis_groups(axis_env, axis_name)
507 def axis_groups(axis_env: AxisEnv, name): function
1288 xla.axis_groups(axis_env, axis_env.names[-1]))