[core] fix autoencoderkl qwenimage for xla

[sayakpaul]

What does this PR do?

Verify with the following:

Code

"""Minimal reproducer: XLA rejects out-of-bounds negative index that CPU/CUDA silently clamps."""

import torch
import torch_xla.core.xla_model as xm

CACHE_T = 2

# Temporal dim = 1 (single image, not video)
x_cpu = torch.randn(1, 4, 1, 8, 8)

# CPU: silently clamps -2 to 0, returns the whole tensor
result_cpu = x_cpu[:, :, -CACHE_T:, :, :]
print(f"CPU: x.shape={x_cpu.shape}, x[:,:,-2:,:,:].shape={result_cpu.shape}")
assert result_cpu.shape == (1, 4, 1, 8, 8), "CPU should clamp and return full tensor"

# CPU with fix: identical result
result_cpu_fix = x_cpu[:, :, -min(CACHE_T, x_cpu.shape[2]):, :, :]
print(f"CPU fix: x[:,:,-min(2,1):,:,:].shape={result_cpu_fix.shape}")
assert torch.equal(result_cpu, result_cpu_fix), "Fix produces same result on CPU"

# XLA: strict bounds checking
device = xm.xla_device()
x_xla = x_cpu.to(device)

print("\nXLA: trying x[:,:,-2:,:,:] with temporal dim=1...")
try:
    result_xla = x_xla[:, :, -CACHE_T:, :, :]
    print(f"XLA: succeeded (shape={result_xla.shape})")
except RuntimeError as e:
    print(f"XLA: FAILED as expected — {e}")

# XLA with fix: works
result_xla_fix = x_xla[:, :, -min(CACHE_T, x_xla.shape[2]):, :, :]
print(f"XLA fix: x[:,:,-min(2,1):,:,:].shape={result_xla_fix.shape}")

# Also verify temporal dim >= 2 works identically with both approaches
print("\n--- Temporal dim = 3 (no issue on either backend) ---")
x_cpu_3 = torch.randn(1, 4, 3, 8, 8)
x_xla_3 = x_cpu_3.to(device)

result_orig = x_xla_3[:, :, -CACHE_T:, :, :]
result_fix = x_xla_3[:, :, -min(CACHE_T, x_xla_3.shape[2]):, :, :]
print(f"Original: shape={result_orig.shape}, Fix: shape={result_fix.shape}")
assert result_orig.shape == result_fix.shape == (1, 4, 2, 8, 8)
assert torch.equal(result_orig.cpu(), result_fix.cpu()), "Fix must produce same values when temporal dim >= CACHE_T"

# Also compare CPU vs XLA values for the fix (temporal dim=1)
assert torch.equal(result_cpu_fix, result_xla_fix.cpu()), "CPU and XLA fix must produce same values"

# And CPU vs XLA for temporal dim >= 2
result_cpu_3_orig = x_cpu_3[:, :, -CACHE_T:, :, :]
assert torch.equal(result_cpu_3_orig, result_orig.cpu()), "CPU and XLA must produce same values when temporal dim >= CACHE_T"

print("\nAll checks passed!")

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[dg845]

Thanks!