jax-ml/jax · error · ValueError

Delta orthogonal initializer requires a 3D, 4D or 5D shape.

Error message

Delta orthogonal initializer requires a 3D, 4D or 5D shape.

What it means

jax.nn.initializers.delta_orthogonal() requires the weight shape to describe a convolutional kernel: 3D (k, fan_in, fan_out), 4D, or 5D. It works by inserting an orthogonal matrix into the center slice of a zero kernel, which only makes sense for conv-style shapes.

Source

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

          [ 0.9120717 ,  0.04322892,  0.40774566],
          [-0.30085585, -0.6050892 ,  0.73712474]],
  <BLANKLINE>
         [[ 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. If you want a dense orthogonal matrix, use jax.nn.initializers.orthogonal() instead
2. For conv kernels, pass the full kernel shape, e.g. (kh, kw, fan_in, fan_out) for conv2d
3. Check that the shape passed is the actual parameter shape, not a flattened size

Example fix

// before
w = jax.nn.initializers.delta_orthogonal()(key, (8, 8))
// after
w = jax.nn.initializers.orthogonal()(key, (8, 8))

def delta_ortho_or_dense(init, key, shape):
    assert len(shape) in (3, 4, 5), f'delta_orthogonal needs 3-5D shape, got {shape}'
    return init(key, shape)

null

null

When it happens

Trigger: Calling delta_orthogonal()(key, shape) with len(shape) not in {3,4,5}, e.g. a 2D matrix shape like (8, 8) or a 6D shape.

Common situations: Using delta_orthogonal where plain orthogonal() was intended (2D dense weights), or passing a flattened/reshaped parameter vector instead of the conv kernel shape.

Related errors

• `fan_in` must be less or equal than `fan_out`.
• scan got `length` argument of {} which disagrees with leadin
• conv_general_dilated batch_group_count must divide lhs batch
• conv_general_dilated rhs output feature dimension size must
• conv_general_dilated window and window_strides must have the

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