Repository: tensorflow
Issue: oneDNN matmul is not supported to have a dynamic dimension
Labels: stat:awaiting tensorflower,type:bug,comp:xla,comp:core

### Issue type

Bug

### Have you reproduced the bug with TensorFlow Nightly?

Yes

### Source

binary

### TensorFlow version

v1.12.1-98433-gb64e3d4ae2e 2.15.0-dev20230812

### Custom code

Yes

### OS platform and distribution

Linux Debian 6.3.11

### Mobile device

_No response_

### Python version

3.11

### Bazel version

_No response_

### GCC/compiler version

_No response_

### CUDA/cuDNN version

_No response_

### GPU model and memory

_No response_

### Current behavior?

The following code works for nightly version 2.15.0-dev20230811, but fails for the version 2.15.0-dev20230812. The issue happens when we compile the function.

### Standalone code to reproduce the issue

```shell
import tensorflow.compat.v2 as tf

class _Slice2Idx:
  """Utility to convert numpy basic slices into TF scatter_nd indices."""

  def __init__(self, tensor):
    self.ranges = [tf.range(d, dtype=tf.int32) for d in tensor.shape]

  def __getitem__(self, slices):
    grid = tf.meshgrid(*(rng[sl] for rng, sl in zip(self.ranges, slices)), indexing='ij')
    return tf.stack(grid, axis=-1)

def no_pivot_ldl(matrix, name='no_pivot_ldl'):
  with tf.name_scope(name) as name:
    matrix = tf.convert_to_tensor(matrix)
    triangular_factor = tf.linalg.band_part(matrix, num_lower=-1, num_upper=0)
    slix = _Slice2Idx(triangular_factor)

    def fn(triangular_factor, i):
      column_tail = triangular_factor[..., i+1:, i]
      x = tf.einsum('...i,...j->...ij', column_tail, column_tail)
      idx = slix[i+1:, i+1:]
      triangular_factor = tf.tensor_scatter_nd_update(triangular_factor, idx, x)
      return triangular_factor

    triangular_factor = tf.foldl(
        fn=fn,
        elems=tf.range(tf.shape(triangular_factor)[-1]),
        initializer=triangular_factor)
    return triangular_factor

inp = tf.Variable([[2., 1.], [1., 2.]])
alt_chol_jit = tf.function(no_pivot_ldl, autograph=False, jit_compile=True)
print(no_pivot_ldl(inp))
alt_chol_jit(inp)
```
```


### Relevant log output

```shell
tf.Tensor(
[[2. 0.]
 [1. 1.]], shape=(2, 2), dtype=float32)
2023-08-23 11:31:54.059605: I tensorflow/compiler/xla/service/service.cc:168] XLA service 0x5db9ec0 initialized for platform Host (this does not guarantee that XLA will be used). Devices:
2023-08-23 11:31:54.059681: I tensorflow/compiler/xla/service/service.cc:176]   StreamExecutor device (0): Host, Default Version
2023-08-23 11:31:54.068228: I tensorflow/compiler/mlir/tensorflow/utils/dump_mlir_util.cc:269] disabling MLIR crash reproducer, set env var `MLIR_CRASH_REPRODUCER_DIRECTORY` to enable.
2023-08-23 11:31:54.093913: W tensorflow/core/framework/op_kernel.cc:1839] OP_REQUIRES failed at xla_ops.cc:625 : UNIMPLEMENTED: CustomCall "__onednn$matmul" is not supported to have a dynamic dimension

tensorflow.python.framework.errors_impl.UnimplementedError: CustomCall "__onednn$matmul" is not supported to have a dynamic dimension [Op:__inference_no_pivot_ldl_264]
```