jax.extend.core.mapped_aval

Contents

jax.extend.core.mapped_aval#

jax.extend.core.mapped_aval(size, axis, aval)[source]#
Parameters:

  • size (AxisSize)

  • aval (AbstractValue)

Return type:

AbstractValue

Contents