Jane Xu 53bdae736d Add capturable single tensor Adamax (#121183) ...

Finishes the work started in https://github.com/pytorch/pytorch/pull/118697. Thanks @MarouaneMaatouk for the attempt, but due to inactivity I have opened this PR for Adamax. Note that the new capturable implementation is much simpler and I've modified the foreach capturable impl--it now calls fewer kernels and is more easily comparable to forloop.

Next steps:
* This PR discovered two bugs: #121178 and #121238.
* Move the now hefty graph optim tests in test_cuda to use OptimInfo.

Pull Request resolved: https://github.com/pytorch/pytorch/pull/121183
Approved by: https://github.com/albanD

2024-03-07 17:57:02 +00:00

..

test_convolution.py

[CUDA][cuDNN][TF32] Misc TF32 updates (#118781)

2024-02-01 15:32:50 +00:00

test_dropout.py

[BE]: Update flake8 to v6.1.0 and fix lints (#116591)

2024-01-03 06:04:44 +00:00

test_embedding.py

Add sparse compressed fake tensor support (#120920)

2024-03-04 14:38:45 +00:00

test_init.py

[dynamo] Skip some tests using scipy.kstest (#116263)

2023-12-21 18:08:29 +00:00

test_lazy_modules.py

Add capturable single tensor Adamax (#121183)

2024-03-07 17:57:02 +00:00

test_load_state_dict.py

Add assign argument to torch.Tensor.module_load (#121158)

2024-03-06 01:32:06 +00:00

test_module_hooks.py

Fix module pre bw hooks when input doesn't req grad but gradients are changed by the user (#116454)

2023-12-28 18:32:50 +00:00

test_multihead_attention.py

[BE]: Update ruff to 0.285 (#107519)

2023-08-22 23:16:38 +00:00

test_packed_sequence.py

[BE]: Apply ruff PERF fixes to torch (#104917)

2023-07-11 20:45:21 +00:00

test_parametrization.py

Spectral norm value test (#121068)

2024-03-05 14:46:31 +00:00

test_pooling.py

Fix edge case for size 1 channels dim in AdaptiveMaxPool (#116482)

2023-12-28 15:02:29 +00:00

test_pruning.py

[BE]: Enable RUF015 codebase wide (#115507)

2023-12-11 15:51:01 +00:00