jax-ml/jax · error · ValueError

`fan_in` must be less or equal than `fan_out`.

Error message

`fan_in` must be less or equal than `fan_out`. 

What it means

delta_orthogonal builds an orthogonal matrix of shape (fan_in, fan_out) = shape[-2:] and requires fan_in <= fan_out (JAX convention: columns are outputs). If the last dimension is smaller than the second-to-last, the QR-based construction cannot be centered in the kernel as required.

Source

Thrown at jax/_src/nn/initializers.py:693

         [[ 0.        ,  0.        ,  0.        ],
          [ 0.        ,  0.        ,  0.        ],
          [ 0.        ,  0.        ,  0.        ]]], dtype=float32)


  .. _delta orthogonal initializer: https://arxiv.org/abs/1806.05393
  """
  def init(key: Array,
           shape: core.Shape,
           dtype: DTypeLikeInexact | None = dtype,
           out_sharding: OutShardingType = None) -> Array:
    if out_sharding is not None:
      raise NotImplementedError
    dtype = dtypes.default_float_dtype() if dtype is None else dtype
    if len(shape) not in [3, 4, 5]:
      raise ValueError("Delta orthogonal initializer requires a 3D, 4D or 5D "
                       "shape.")
    if shape[-1] < shape[-2]:
      raise ValueError("`fan_in` must be less or equal than `fan_out`. ")
    ortho_init = orthogonal(scale=scale, column_axis=column_axis, dtype=dtype)
    ortho_matrix = ortho_init(key, shape[-2:])
    W = jnp.zeros(shape, dtype=dtype)
    if len(shape) == 3:
      k = shape[0]
      return W.at[(k-1)//2, ...].set(ortho_matrix)
    elif len(shape) == 4:
      k1, k2 = shape[:2]
      return W.at[(k1-1)//2, (k2-1)//2, ...].set(ortho_matrix)
    else:
      k1, k2, k3 = shape[:3]
      return W.at[(k1-1)//2, (k2-1)//2, (k3-1)//2, ...].set(ortho_matrix)
  return init

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

  1. Transpose the last two dims of your kernel layout so fan_out >= fan_in (swap channel order in the layer definition)
  2. Use orthogonal() or a variance-scaling initializer for channel-reducing layers
  3. If channel reduction is required, initialize on the transposed shape and transpose the result back

Example fix

// before
w = jax.nn.initializers.delta_orthogonal()(key, (3, 64, 32))  # fan_out < fan_in
// after
w = jax.nn.initializers.delta_orthogonal()(key, (3, 32, 64)).transpose(0, 2, 1)
Defensive patterns

Strategy: validation

Validate before calling

def delta_ortho_checked(init, key, shape):
    assert shape[-1] >= shape[-2], 'fan_out (shape[-1]) must be >= fan_in (shape[-2])'
    return init(key, shape)

Type guard

null

Try / catch

null

Prevention

When it happens

Trigger: Calling delta_orthogonal()(key, shape) where shape[-1] < shape[-2], e.g. (3, 64, 32) for a channel-increasing convolution.

Common situations: Conv layers that reduce channels (e.g. 64 input channels to 32 output channels), or using a different fan_in/fan_out axis convention than JAX expects.

Related errors


AI-assisted analysis of jax-ml/jax@1e1c6a8fc0 (2026-08-27). Data as JSON: /api/errors/366510463ac70bf2. Report an issue: GitHub.