def _GetFans(shape, out_dim=-1, in_dim=-2, nonreceptive_dims=None):
  """Get the fan-in and fan-out sizes for the given shape and dims."""
  # Temporary fix until numpy.delete supports negative indices.
  if out_dim < 0:
    out_dim += len(shape)
  if in_dim < 0:
    in_dim += len(shape)

  if nonreceptive_dims is None:
    nonreceptive_dims = []
  if not isinstance(nonreceptive_dims, (list, tuple)):
    nonreceptive_dims = [nonreceptive_dims]

  receptive_field = jnp.prod(np.delete(shape, [in_dim, out_dim,
                                               *nonreceptive_dims]))
  if len(shape) >= 2:
    fan_in, fan_out = shape[in_dim], shape[out_dim]
  elif len(shape) == 1:
    fan_in = shape[0]
    fan_out = shape[0]
  else:
    fan_in = 1.
    fan_out = 1.
    fan_in *= receptive_field
    fan_out *= receptive_field
  return fan_in, fan_out
Get the fan-in and fan-out sizes for the given shape and dims.
trax/layers/initializers.py