Repository: pytorch
Issue: flex_attention + `export_for_inference` ` NYI: querying is_contiguous inside of vmap for memory_format other than torch.contiguous_format`
Labels: triaged,oncall: pt2,module: higher order operators,module: pt2-dispatcher,dynamo-functorch,module: flex attention

### 🐛 Describe the bug

```python
import torch
from torch import nn, Tensor
from torch.export import export_for_inference, Dim
from torch.nn.attention.flex_attention import flex_attention

class MinimalAttention(nn.Module):
    """Minimal reproduction of Attention layer using flex_attention."""
    
    def __init__(self, dim, num_heads=8, qkv_bias=True):
        super().__init__()
        self.num_heads = num_heads
        self.head_dim = dim // num_heads
        self.scale = self.head_dim ** -0.5

        self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
        self.proj = nn.Linear(dim, dim)

    def forward(self, x: Tensor) -> Tensor:
        B, N, C = x.shape  # Flattened spatial dimensions with batch
        qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, self.head_dim).permute(0, 3, 2, 1, 4)
        q, k, v = qkv[:, :, 0], qkv[:, :, 1], qkv[:, :, 2]

        # Compute attention output using flex_attention
        attn_out = flex_attention(q, k, v, scale=self.scale)
        return self.proj(attn_out.transpose(1, 2).reshape(B, N, C))

def replicate_error():
    # Define the input tensor with fixed values to avoid randomness
    B, N, C = 2, 1024, 64
    x = torch.ones(B, N, C, device="cuda", dtype=torch.float32)  # Fixed values instead of random

    # Initialize the minimal attention model
    model = MinimalAttention(dim=C).cuda()

    # Set up dynamic shape for batch size flexibility
    dynamic_shapes = {
        "x": {0: Dim("batch", min=1, max=B), 1: Dim.STATIC, 2: Dim.STATIC}
    }

    # Export with AMP and dynamic shape settings
    try:
        with torch.amp.autocast(device_type="cuda", dtype=torch.float16):
            exported_model = export_for_inference(model, (x,), dynamic_shapes=dynamic_shapes)
        torch.save(exported_model, "minimal_attention_error_model.pt")
        print("Exported model successfully.")
    except RuntimeError as e:
        print("Export failed with error:", e)

# Run the minimal reproduction
replicate_error()
```

```python
Export failed with error: NYI: querying is_contiguous inside of vmap for memory_format other than torch.contiguous_format

While executing %dense_mask : [num_users=2] = call_method[target=new_zeros](args = (%_add_batch_dim_3, 1, 2), kwargs = {dtype: torch.int32})
Original traceback:
  File "/workspace/delme_repro.py", line 24, in forward
    attn_out = flex_attention(q, k, v, scale=self.scale)
  File "/opt/conda/lib/python3.11/site-packages/torch/nn/attention/flex_attention.py", line 1197, in flex_attention
    block_mask = _create_empty_block_mask(query, key)
  File "/opt/conda/lib/python3.11/site-packages/torch/nn/attention/flex_attention.py", line 893, in _create_empty_block_mask
    return BlockMask.from_kv_blocks(
  File "/opt/conda/lib/python3.11/site-packages/torch/nn/attention/flex_attention.py", line 349, in from_kv_blocks
    q_num_blocks, q_indices = _transpose_ordered(kv_num_blocks, kv_indices)
  File "/opt/conda/lib/python3.11/site-packages/torch/nn/attention/flex_attention.py", line 187, in _transpose_ordered
    dense = _ordered_to_dense(num_blocks_in_row, col_indices)
  File "/opt/conda/lib/python3.11/site-packages/torch/nn/attention/flex_attention.py", line 172, in _ordered_to_dense
    out = create_dense_batched(num_blocks_in_row, col_indices)
  File "/opt/conda/lib/python3.11/site-packages/torch/_functorch/apis.py", line 203, in wrapped
    return vmap_impl(
  File "/opt/conda/lib/python3.11/site-packages/torch/_functorch/vmap.py", line 331, in vmap_impl
    return _flat_vmap(
  File "/opt/conda/lib/python3.11/site-packages/torch/_functorch/vmap.py", line 479, in _flat_vmap
    batched_outputs = func(*batched_inputs, **kwargs)
  File "/opt/conda/lib/python3.11/site-packages/torch/_functorch/apis.py", line 203, in wrapped
    return vmap_impl(
  File "/opt/conda/lib/python3.11/site-packages/torch/_functorch/vmap.py", line 331, in vmap_impl
    return _flat_vmap(
  File "/opt/conda/lib/python3.11/site-packages/torch/_functorch/vmap.py", line 479, in _flat_vmap
    batched_outputs = func(*batched_inputs, **kwargs)
  File "/opt/conda/lib/python3.11/site-packages/torch/nn/attention/flex_attention.py", line 153, in create_dense_one
    dense_mask = kv_indices.new_zeros(num_rows, num_cols + 1, dtype=torch.int32)
 ```
 
 /cc @chauhang @penguinwu @zou3519 @ydwu4 @bdhirsh @yf225 @Chillee @drisspg @yanboliang @BoyuanFeng @ezyang @tugsbayasgalan 

### Versions

nightly