jax.extend.core.unmapped_aval

Contents

jax.extend.core.unmapped_aval#

jax.extend.core.unmapped_aval(size, axis, aval, explicit_mesh_axis=None)[source]#
Parameters:

  • size (AxisSize)

  • axis (int | None)

  • aval (AbstractValue)

Return type:

AbstractValue

Contents