e164c3574da160fb72fc962947a783fa299424e3,sonnet/src/batch_apply.py,,merge_leading_dims,#Any#Any#,97

Before Change

  if len(x.shape) < num_dims:
    return x
  return tf.reshape(x, [-1] + x.shape.as_list()[num_dims:])

After Change

    return x

  // Check if the rank of the input tensor is well-defined.
  if x.shape.dims is None:
    raise ValueError(
        "Can"t merge leading dimensions of tensor of unknown rank.")

  // We can only merge the num_dims leading dimensions if the rank of the given
  // tensor is sufficiently large.
  if num_dims > x.shape.rank:
    return x

  static_shape = [maybe_prod(x.shape[:num_dims])] + x.shape[num_dims:]

Instances