import torch
import torch.autograd.functional as AF

x = torch.randn(16, 4, 16, dtype=torch.float64, requires_grad=True)

# [A] forward-mode with vectorize=False (default) — fails at execution time
print("[A] strategy='forward-mode'  vectorize=False (default)")
try:
    J = AF.jacobian(torch.tanh, x, strategy='forward-mode', vectorize=False)
    print(f"  ok: {tuple(J.shape)}")
except Exception as e:
    print(f"  {type(e).__name__}: {e}")

# [B] forward-mode with vectorize=True — the required combination, works
print("[B] strategy='forward-mode'  vectorize=True")
try:
    J = AF.jacobian(torch.tanh, x, strategy='forward-mode', vectorize=True)
    print(f"  ok: {tuple(J.shape)}")
except Exception as e:
    print(f"  {type(e).__name__}: {e}")

# [C] reverse-mode with vectorize=False — the default, works fine
x2 = torch.randn(4, 4, dtype=torch.float64, requires_grad=True)
print("[C] strategy='reverse-mode'  vectorize=False (default)")
try:
    J = AF.jacobian(torch.tanh, x2, strategy='reverse-mode', vectorize=False)
    print(f"  ok: {tuple(J.shape)}")
except Exception as e:
    print(f"  {type(e).__name__}: {e}")

---

[A] strategy='forward-mode'  vectorize=False (default)
  NotImplementedError: Computing Jacobian using forward-AD or forward-over-reverse
  Hessian is only implemented for `vectorize=True`.
[B] strategy='forward-mode'  vectorize=True
  ok: (16, 4, 16, 16, 4, 16)
[C] strategy='reverse-mode'  vectorize=False (default)
  ok: (4, 4, 4, 4)

---

[A] strategy='forward-mode'  vectorize=False (default)
  ValueError: strategy='forward-mode' requires vectorize=True.
  Pass vectorize=True or use strategy='reverse-mode'.
[B] strategy='forward-mode'  vectorize=True
  ok: (16, 4, 16, 16, 4, 16)
[C] strategy='reverse-mode'  vectorize=False (default)
  ok: (4, 4, 4, 4)

---

if strategy in ('forward-mode', 'forward-over-reverse') and not vectorize:
    raise ValueError(
        f"strategy='{strategy}' requires vectorize=True."
    )

---

PyTorch version: 2.13.0.dev20260512+cu130
Is debug build: False
CUDA used to build PyTorch: 13.0
ROCM used to build PyTorch: N/A

OS: Ubuntu 24.04.4 LTS (x86_64)
GCC version: (Ubuntu 13.3.0-6ubuntu2~24.04.1) 13.3.0
Clang version: 18.1.3 (1ubuntu1)
CMake version: version 3.28.3
Libc version: glibc-2.39

Python version: 3.12.3 (main, Mar 23 2026, 19:04:32) [GCC 13.3.0] (64-bit runtime)
Python platform: Linux-6.14.0-37-generic-x86_64-with-glibc2.39
Is CUDA available: True
GPU models and configuration: GPU 0: NVIDIA GeForce RTX 5090
Nvidia driver version: 590.48.01

[pip3] numpy==2.4.4
[pip3] torch==2.13.0.dev20260512+cu130
[pip3] triton==3.7.0+git88b227e2