jax
jax copied to clipboard
google
Stochastic jaxlib.xla_extension.XlaRuntimeError with sine activation function on NVIDIA A100 (all algorithms tried for... failed.)
Description
I sometimes, but not consistently, get the following jaxlib.xla_extension.XlaRuntimeError when training a neural network with sine activations in JAX on an NVIDIA A100 GPU:
jaxlib.xla_extension.XlaRuntimeError: INTERNAL: All algorithms tried for %dot.128 = f32[384,384]{1,0} dot(f32[216000,384]{1,0} %multiply.51, f32[216000,384]{1,0} %sine.3), lhs_contracting_dims={0}, rhs_contracting_dims={0}, metadata={op_name="jit(train_step)/jit(main)/dot_general[dimension_numbers=(((1,), (0,)), ((), ())) precision=None preferred_element_type=float32]" source_file="./layers.py" source_line=160} failed. Falling back to default algorithm. Per-algorithm errors:
Results do not match the reference. This is likely a bug/unexpected loss of precision.
...
Results do not match the reference. This is likely a bug/unexpected loss of precision.
The final part of the traceback is
File "/home/skoop/.conda/envs/crystals/lib/python3.10/site-packages/equinox/_jit.py", line 206, in __call__
return self._call(False, args, kwargs)
File "/home/skoop/.conda/envs/crystals/lib/python3.10/site-packages/equinox/_module.py", line 935, in __call__
return self.__func__(self.__self__, *args, **kwargs)
File "/home/skoop/.conda/envs/crystals/lib/python3.10/site-packages/equinox/_jit.py", line 200, in _call
out = self._cached(dynamic_donate, dynamic_nodonate, static)
This error does however not occur when I use an NVIDIA GeForce RTX 3080 or NVIDIA GeForce RTX 3060.
Some accompanying logs are
2024-02-16 16:05:48.264495: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 3557: -12.5601, expected -8.76762
2024-02-16 16:05:48.264583: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 7281: 17.2671, expected 26.1401
2024-02-16 16:05:48.264685: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 31490: 46.9089, expected 35.7881
2024-02-16 16:05:48.264734: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 42098: 92.8985, expected 82.6593
2024-02-16 16:05:48.264871: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 77426: 14.7896, expected 10.7041
2024-02-16 16:05:48.264960: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 99275: 61.807, expected 51.012
2024-02-16 16:05:48.264979: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 101673: 54.8168, expected 64.0151
2024-02-16 16:05:48.265027: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 111729: -23.3462, expected -28.2862
2024-02-16 16:05:48.265046: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 114033: -54.0808, expected -44.5037
2024-02-16 16:05:48.265094: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 124619: -35.2514, expected -46.6975
2024-02-16 16:05:48.265104: E external/xla/xla/service/gpu/triton_autotuner.cc:816] Results do not match the reference. This is likely a bug/unexpected loss of precision.
2024-02-16 16:05:48.297565: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 3557: -12.5601, expected -8.76762
2024-02-16 16:05:48.297609: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 7281: 17.2671, expected 26.1401
2024-02-16 16:05:48.297709: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 31490: 46.9089, expected 35.7881
2024-02-16 16:05:48.297757: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 42098: 92.8985, expected 82.6593
2024-02-16 16:05:48.297901: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 77426: 14.7896, expected 10.7041
2024-02-16 16:05:48.297993: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 99275: 61.807, expected 51.012
2024-02-16 16:05:48.298012: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 101673: 54.8168, expected 64.0151
2024-02-16 16:05:48.298058: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 111729: -23.3462, expected -28.2862
2024-02-16 16:05:48.298076: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 114033: -54.0808, expected -44.5037
2024-02-16 16:05:48.298124: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 124619: -35.2514, expected -46.6975
2024-02-16 16:05:48.298133: E external/xla/xla/service/gpu/triton_autotuner.cc:816] Results do not match the reference. This is likely a bug/unexpected loss of precision.
2024-02-16 16:05:48.933430: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 3557: -12.5601, expected -8.76762
2024-02-16 16:05:48.933503: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 7281: 17.2671, expected 26.1401
2024-02-16 16:05:48.933612: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 31490: 46.9089, expected 35.7881
2024-02-16 16:05:48.933663: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 42098: 92.8985, expected 82.6593
2024-02-16 16:05:48.933808: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 77426: 14.7896, expected 10.7041
2024-02-16 16:05:48.933897: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 99275: 61.807, expected 51.012
2024-02-16 16:05:48.933916: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 101673: 54.8168, expected 64.0151
2024-02-16 16:05:48.933962: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 111729: -23.3462, expected -28.2862
2024-02-16 16:05:48.933981: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 114033: -54.0808, expected -44.5037
2024-02-16 16:05:48.934031: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 124619: -35.2514, expected -46.6975
2024-02-16 16:05:48.934042: E external/xla/xla/service/gpu/triton_autotuner.cc:816] Results do not match the reference. This is likely a bug/unexpected loss of precision.
2024-02-16 16:05:49.506808: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 3557: -12.5601, expected -8.76762
2024-02-16 16:05:49.506871: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 7281: 17.2671, expected 26.1401
2024-02-16 16:05:49.506970: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 31490: 46.9089, expected 35.7881
2024-02-16 16:05:49.507018: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 42098: 92.8985, expected 82.6593
2024-02-16 16:05:49.507155: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 77426: 14.7896, expected 10.7041
2024-02-16 16:05:49.507243: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 99275: 61.807, expected 51.012
2024-02-16 16:05:49.507262: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 101673: 54.8168, expected 64.0151
2024-02-16 16:05:49.507308: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 111729: -23.3462, expected -28.2862
2024-02-16 16:05:49.507326: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 114033: -54.0808, expected -44.5037
2024-02-16 16:05:49.507374: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 124619: -35.2514, expected -46.6975
2024-02-16 16:05:49.507384: E external/xla/xla/service/gpu/triton_autotuner.cc:816] Results do not match the reference. This is likely a bug/unexpected loss of precision.
2024-02-16 16:05:51.473397: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 3557: -12.4619, expected -8.76762
2024-02-16 16:05:51.473461: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 7281: 17.3246, expected 26.1401
2024-02-16 16:05:51.473567: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 31490: 46.9092, expected 35.7881
2024-02-16 16:05:51.473616: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 42098: 92.9825, expected 82.6593
2024-02-16 16:05:51.473762: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 77426: 14.8192, expected 10.7041
2024-02-16 16:05:51.473796: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 84145: 0.918945, expected 1.13448
2024-02-16 16:05:51.473861: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 99275: 61.8836, expected 51.012
2024-02-16 16:05:51.473880: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 101673: 54.8101, expected 64.0151
2024-02-16 16:05:51.473926: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 111729: -23.3323, expected -28.2862
2024-02-16 16:05:51.473945: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 114033: -54.0575, expected -44.5037
2024-02-16 16:05:51.473955: E external/xla/xla/service/gpu/triton_autotuner.cc:816] Results do not match the reference. This is likely a bug/unexpected loss of precision.
2024-02-16 16:05:51.502365: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 3557: -12.449, expected -8.76762
2024-02-16 16:05:51.502411: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 7281: 17.3652, expected 26.1401
2024-02-16 16:05:51.502515: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 31490: 46.9354, expected 35.7881
2024-02-16 16:05:51.502574: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 42098: 92.9954, expected 82.6593
2024-02-16 16:05:51.502723: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 77426: 14.8499, expected 10.7041
2024-02-16 16:05:51.502812: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 99275: 61.8535, expected 51.012
2024-02-16 16:05:51.502831: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 101673: 54.8031, expected 64.0151
2024-02-16 16:05:51.502877: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 111729: -23.2869, expected -28.2862
2024-02-16 16:05:51.502896: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 114033: -53.988, expected -44.5037
2024-02-16 16:05:51.502944: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 124619: -35.3679, expected -46.6975
2024-02-16 16:05:51.502954: E external/xla/xla/service/gpu/triton_autotuner.cc:816] Results do not match the reference. This is likely a bug/unexpected loss of precision.
2024-02-16 16:05:51.531594: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 3557: -12.4182, expected -8.76762
2024-02-16 16:05:51.531638: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 7281: 17.3456, expected 26.1401
2024-02-16 16:05:51.531741: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 31490: 46.8993, expected 35.7881
2024-02-16 16:05:51.531761: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 34034: 64.7217, expected 58.1465
2024-02-16 16:05:51.531802: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 42098: 93.0271, expected 82.6593
2024-02-16 16:05:51.531940: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 77426: 14.8545, expected 10.7041
2024-02-16 16:05:51.532029: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 99275: 61.8445, expected 51.012
2024-02-16 16:05:51.532048: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 101673: 54.8469, expected 64.0151
2024-02-16 16:05:51.532094: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 111729: -23.2837, expected -28.2862
2024-02-16 16:05:51.532113: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 114033: -53.994, expected -44.5037
2024-02-16 16:05:51.532122: E external/xla/xla/service/gpu/triton_autotuner.cc:816] Results do not match the reference. This is likely a bug/unexpected loss of precision.
2024-02-16 16:05:52.266140: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 3557: -12.5601, expected -8.76762
2024-02-16 16:05:52.266194: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 7281: 17.2671, expected 26.1401
2024-02-16 16:05:52.266300: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 31490: 46.9089, expected 35.7881
2024-02-16 16:05:52.266349: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 42098: 92.8985, expected 82.6593
2024-02-16 16:05:52.266492: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 77426: 14.7896, expected 10.7041
2024-02-16 16:05:52.266589: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 99275: 61.807, expected 51.012
2024-02-16 16:05:52.266609: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 101673: 54.8168, expected 64.0151
2024-02-16 16:05:52.266656: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 111729: -23.3462, expected -28.2862
2024-02-16 16:05:52.266676: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 114033: -54.0808, expected -44.5037
2024-02-16 16:05:52.266724: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 124619: -35.2514, expected -46.6975
2024-02-16 16:05:52.266734: E external/xla/xla/service/gpu/triton_autotuner.cc:816] Results do not match the reference. This is likely a bug/unexpected loss of precision.
2024-02-16 16:05:52.889293: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 3557: -12.5601, expected -8.76762
2024-02-16 16:05:52.889346: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 7281: 17.2671, expected 26.1401
2024-02-16 16:05:52.889456: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 31490: 46.9089, expected 35.7881
2024-02-16 16:05:52.889505: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 42098: 92.8985, expected 82.6593
2024-02-16 16:05:52.889651: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 77426: 14.7896, expected 10.7041
2024-02-16 16:05:52.889743: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 99275: 61.807, expected 51.012
2024-02-16 16:05:52.889762: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 101673: 54.8168, expected 64.0151
2024-02-16 16:05:52.889808: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 111729: -23.3462, expected -28.2862
2024-02-16 16:05:52.889826: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 114033: -54.0808, expected -44.5037
2024-02-16 16:05:52.889874: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 124619: -35.2514, expected -46.6975
2024-02-16 16:05:52.889884: E external/xla/xla/service/gpu/triton_autotuner.cc:816] Results do not match the reference. This is likely a bug/unexpected loss of precision.
2024-02-16 16:05:52.909936: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 3557: -12.4314, expected -8.76762
2024-02-16 16:05:52.909981: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 7281: 17.3356, expected 26.1401
2024-02-16 16:05:52.910086: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 31490: 46.9041, expected 35.7881
2024-02-16 16:05:52.910134: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 42098: 92.9747, expected 82.6593
2024-02-16 16:05:52.910272: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 77426: 14.8142, expected 10.7041
2024-02-16 16:05:52.910361: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 99275: 61.8406, expected 51.012
2024-02-16 16:05:52.910380: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 101673: 54.7939, expected 64.0151
2024-02-16 16:05:52.910427: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 111729: -23.2614, expected -28.2862
2024-02-16 16:05:52.910445: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 114033: -54.0424, expected -44.5037
2024-02-16 16:05:52.910495: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 124619: -35.2993, expected -46.6975
2024-02-16 16:05:52.910505: E external/xla/xla/service/gpu/triton_autotuner.cc:816] Results do not match the reference. This is likely a bug/unexpected loss of precision.
2024-02-16 16:05:52.983981: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 3557: -12.5601, expected -8.76762
2024-02-16 16:05:52.984024: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 7281: 17.2671, expected 26.1401
2024-02-16 16:05:52.984129: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 31490: 46.9089, expected 35.7881
2024-02-16 16:05:52.984179: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 42098: 92.8985, expected 82.6593
2024-02-16 16:05:52.984323: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 77426: 14.7896, expected 10.7041
2024-02-16 16:05:52.984412: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 99275: 61.807, expected 51.012
2024-02-16 16:05:52.984431: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 101673: 54.8168, expected 64.0151
2024-02-16 16:05:52.984479: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 111729: -23.3462, expected -28.2862
2024-02-16 16:05:52.984497: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 114033: -54.0808, expected -44.5037
2024-02-16 16:05:52.984546: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 124619: -35.2514, expected -46.6975
2024-02-16 16:05:52.984561: E external/xla/xla/service/gpu/triton_autotuner.cc:816] Results do not match the reference. This is likely a bug/unexpected loss of precision.
2024-02-16 16:05:52.990249: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 3557: -12.4182, expected -8.76762
2024-02-16 16:05:52.990290: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 7281: 17.3456, expected 26.1401
2024-02-16 16:05:52.990396: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 31490: 46.8993, expected 35.7881
2024-02-16 16:05:52.990415: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 34034: 64.7217, expected 58.1465
2024-02-16 16:05:52.990457: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 42098: 93.0271, expected 82.6593
2024-02-16 16:05:52.990601: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 77426: 14.8545, expected 10.7041
2024-02-16 16:05:52.990696: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 99275: 61.8445, expected 51.012
2024-02-16 16:05:52.990715: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 101673: 54.8469, expected 64.0151
2024-02-16 16:05:52.990761: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 111729: -23.2837, expected -28.2862
2024-02-16 16:05:52.990780: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 114033: -53.994, expected -44.5037
2024-02-16 16:05:52.990789: E external/xla/xla/service/gpu/triton_autotuner.cc:816] Results do not match the reference. This is likely a bug/unexpected loss of precision.
2024-02-16 16:05:53.003926: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 3557: -12.5601, expected -8.76762
2024-02-16 16:05:53.003967: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 7281: 17.2671, expected 26.1401
2024-02-16 16:05:53.004071: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 31490: 46.9089, expected 35.7881
2024-02-16 16:05:53.004122: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 42098: 92.8985, expected 82.6593
2024-02-16 16:05:53.004260: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 77426: 14.7896, expected 10.7041
2024-02-16 16:05:53.004349: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 99275: 61.807, expected 51.012
2024-02-16 16:05:53.004368: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 101673: 54.8168, expected 64.0151
2024-02-16 16:05:53.004414: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 111729: -23.3462, expected -28.2862
2024-02-16 16:05:53.004433: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 114033: -54.0808, expected -44.5037
2024-02-16 16:05:53.004481: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 124619: -35.2514, expected -46.6975
2024-02-16 16:05:53.004491: E external/xla/xla/service/gpu/triton_autotuner.cc:816] Results do not match the reference. This is likely a bug/unexpected loss of precision.
2024-02-16 16:05:53.021304: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 3557: -12.5601, expected -8.76762
2024-02-16 16:05:53.021347: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 7281: 17.2671, expected 26.1401
2024-02-16 16:05:53.021452: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 31490: 46.9089, expected 35.7881
2024-02-16 16:05:53.021503: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 42098: 92.8985, expected 82.6593
2024-02-16 16:05:53.021653: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 77426: 14.7896, expected 10.7041
2024-02-16 16:05:53.021744: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 99275: 61.807, expected 51.012
2024-02-16 16:05:53.021763: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 101673: 54.8168, expected 64.0151
2024-02-16 16:05:53.021811: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 111729: -23.3462, expected -28.2862
2024-02-16 16:05:53.021830: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 114033: -54.0808, expected -44.5037
2024-02-16 16:05:53.021880: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 124619: -35.2514, expected -46.6975
2024-02-16 16:05:53.021890: E external/xla/xla/service/gpu/triton_autotuner.cc:816] Results do not match the reference. This is likely a bug/unexpected loss of precision.
2024-02-16 16:05:53.060989: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 3557: -12.5601, expected -8.76762
2024-02-16 16:05:53.061030: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 7281: 17.2671, expected 26.1401
2024-02-16 16:05:53.061134: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 31490: 46.9089, expected 35.7881
2024-02-16 16:05:53.061185: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 42098: 92.8985, expected 82.6593
2024-02-16 16:05:53.061322: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 77426: 14.7896, expected 10.7041
2024-02-16 16:05:53.061415: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 99275: 61.807, expected 51.012
2024-02-16 16:05:53.061434: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 101673: 54.8168, expected 64.0151
2024-02-16 16:05:53.061480: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 111729: -23.3462, expected -28.2862
2024-02-16 16:05:53.061499: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 114033: -54.0808, expected -44.5037
2024-02-16 16:05:53.061547: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 124619: -35.2514, expected -46.6975
2024-02-16 16:05:53.061561: E external/xla/xla/service/gpu/triton_autotuner.cc:816] Results do not match the reference. This is likely a bug/unexpected loss of precision.
2024-02-16 16:05:53.139898: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 3557: -12.5601, expected -8.76762
2024-02-16 16:05:53.139939: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 7281: 17.2671, expected 26.1401
2024-02-16 16:05:53.140042: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 31490: 46.9089, expected 35.7881
2024-02-16 16:05:53.140092: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 42098: 92.8985, expected 82.6593
2024-02-16 16:05:53.140230: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 77426: 14.7896, expected 10.7041
2024-02-16 16:05:53.140319: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 99275: 61.807, expected 51.012
2024-02-16 16:05:53.140338: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 101673: 54.8168, expected 64.0151
2024-02-16 16:05:53.140384: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 111729: -23.3462, expected -28.2862
2024-02-16 16:05:53.140403: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 114033: -54.0808, expected -44.5037
2024-02-16 16:05:53.140451: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 124619: -35.2514, expected -46.6975
2024-02-16 16:05:53.140461: E external/xla/xla/service/gpu/triton_autotuner.cc:816] Results do not match the reference. This is likely a bug/unexpected loss of precision.
2024-02-16 16:05:53.859318: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 3557: -12.5601, expected -8.76762
2024-02-16 16:05:53.859371: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 7281: 17.2671, expected 26.1401
2024-02-16 16:05:53.859473: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 31490: 46.9089, expected 35.7881
2024-02-16 16:05:53.859524: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 42098: 92.8985, expected 82.6593
2024-02-16 16:05:53.859673: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 77426: 14.7896, expected 10.7041
2024-02-16 16:05:53.859764: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 99275: 61.807, expected 51.012
2024-02-16 16:05:53.859783: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 101673: 54.8168, expected 64.0151
2024-02-16 16:05:53.859829: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 111729: -23.3462, expected -28.2862
2024-02-16 16:05:53.859848: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 114033: -54.0808, expected -44.5037
2024-02-16 16:05:53.859896: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 124619: -35.2514, expected -46.6975
2024-02-16 16:05:53.859906: E external/xla/xla/service/gpu/triton_autotuner.cc:816] Results do not match the reference. This is likely a bug/unexpected loss of precision.
2024-02-16 16:05:54.787867: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 3557: -12.5601, expected -8.76762
2024-02-16 16:05:54.787922: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 7281: 17.2671, expected 26.1401
2024-02-16 16:05:54.788026: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 31490: 46.9089, expected 35.7881
2024-02-16 16:05:54.788074: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 42098: 92.8985, expected 82.6593
2024-02-16 16:05:54.788212: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 77426: 14.7896, expected 10.7041
2024-02-16 16:05:54.788305: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 99275: 61.807, expected 51.012
2024-02-16 16:05:54.788323: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 101673: 54.8168, expected 64.0151
2024-02-16 16:05:54.788369: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 111729: -23.3462, expected -28.2862
2024-02-16 16:05:54.788388: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 114033: -54.0808, expected -44.5037
2024-02-16 16:05:54.788438: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 124619: -35.2514, expected -46.6975
2024-02-16 16:05:54.788448: E external/xla/xla/service/gpu/triton_autotuner.cc:816] Results do not match the reference. This is likely a bug/unexpected loss of precision.
2024-02-16 16:05:54.808215: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 3557: -12.5601, expected -8.76762
2024-02-16 16:05:54.808259: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 7281: 17.2671, expected 26.1401
2024-02-16 16:05:54.808362: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 31490: 46.9089, expected 35.7881
2024-02-16 16:05:54.808413: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 42098: 92.8985, expected 82.6593
2024-02-16 16:05:54.808551: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 77426: 14.7896, expected 10.7041
2024-02-16 16:05:54.808652: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 99275: 61.807, expected 51.012
2024-02-16 16:05:54.808672: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 101673: 54.8168, expected 64.0151
2024-02-16 16:05:54.808720: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 111729: -23.3462, expected -28.2862
2024-02-16 16:05:54.808739: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 114033: -54.0808, expected -44.5037
2024-02-16 16:05:54.808787: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 124619: -35.2514, expected -46.6975
2024-02-16 16:05:54.808796: E external/xla/xla/service/gpu/triton_autotuner.cc:816] Results do not match the reference. This is likely a bug/unexpected loss of precision.
2024-02-16 16:05:55.033790: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 3557: -12.5601, expected -8.76762
2024-02-16 16:05:55.033834: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 7281: 17.2671, expected 26.1401
2024-02-16 16:05:55.033938: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 31490: 46.9089, expected 35.7881
2024-02-16 16:05:55.033990: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 42098: 92.8985, expected 82.6593
2024-02-16 16:05:55.034134: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 77426: 14.7896, expected 10.7041
2024-02-16 16:05:55.034223: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 99275: 61.807, expected 51.012
2024-02-16 16:05:55.034242: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 101673: 54.8168, expected 64.0151
2024-02-16 16:05:55.034290: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 111729: -23.3462, expected -28.2862
2024-02-16 16:05:55.034309: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 114033: -54.0808, expected -44.5037
2024-02-16 16:05:55.034357: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 124619: -35.2514, expected -46.6975
2024-02-16 16:05:55.034367: E external/xla/xla/service/gpu/triton_autotuner.cc:816] Results do not match the reference. This is likely a bug/unexpected loss of precision.
2024-02-16 16:05:55.078042: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 3557: -12.5601, expected -8.76762
2024-02-16 16:05:55.078084: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 7281: 17.2671, expected 26.1401
2024-02-16 16:05:55.078190: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 31490: 46.9089, expected 35.7881
2024-02-16 16:05:55.078240: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 42098: 92.8985, expected 82.6593
2024-02-16 16:05:55.078378: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 77426: 14.7896, expected 10.7041
2024-02-16 16:05:55.078468: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 99275: 61.807, expected 51.012
2024-02-16 16:05:55.078487: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 101673: 54.8168, expected 64.0151
2024-02-16 16:05:55.078533: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 111729: -23.3462, expected -28.2862
2024-02-16 16:05:55.078551: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 114033: -54.0808, expected -44.5037
2024-02-16 16:05:55.078607: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 124619: -35.2514, expected -46.6975
2024-02-16 16:05:55.078616: E external/xla/xla/service/gpu/triton_autotuner.cc:816] Results do not match the reference. This is likely a bug/unexpected loss of precision.
2024-02-16 16:05:55.551993: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 3557: -12.5601, expected -8.76762
2024-02-16 16:05:55.552042: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 7281: 17.2671, expected 26.1401
2024-02-16 16:05:55.552145: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 31490: 46.9089, expected 35.7881
2024-02-16 16:05:55.552198: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 42098: 92.8985, expected 82.6593
2024-02-16 16:05:55.552341: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 77426: 14.7896, expected 10.7041
2024-02-16 16:05:55.552430: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 99275: 61.807, expected 51.012
2024-02-16 16:05:55.552449: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 101673: 54.8168, expected 64.0151
2024-02-16 16:05:55.552497: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 111729: -23.3462, expected -28.2862
2024-02-16 16:05:55.552515: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 114033: -54.0808, expected -44.5037
2024-02-16 16:05:55.552570: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 124619: -35.2514, expected -46.6975
2024-02-16 16:05:55.552579: E external/xla/xla/service/gpu/triton_autotuner.cc:816] Results do not match the reference. This is likely a bug/unexpected loss of precision.
2024-02-16 16:05:55.930834: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 3557: -12.4384, expected -8.76762
2024-02-16 16:05:55.930882: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 7281: 17.2967, expected 26.1401
2024-02-16 16:05:55.930988: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 31490: 46.8611, expected 35.7881
2024-02-16 16:05:55.931036: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 42098: 92.9786, expected 82.6593
2024-02-16 16:05:55.931174: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 77426: 14.8248, expected 10.7041
2024-02-16 16:05:55.931266: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 99275: 61.8658, expected 51.012
2024-02-16 16:05:55.931285: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 101673: 54.8564, expected 64.0151
2024-02-16 16:05:55.931333: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 111729: -23.2874, expected -28.2862
2024-02-16 16:05:55.931352: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 114033: -54.0519, expected -44.5037
2024-02-16 16:05:55.931401: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 124619: -35.3099, expected -46.6975
2024-02-16 16:05:55.931411: E external/xla/xla/service/gpu/triton_autotuner.cc:816] Results do not match the reference. This is likely a bug/unexpected loss of precision.
2024-02-16 16:05:56.905503: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 3557: -12.4314, expected -8.76762
2024-02-16 16:05:56.905563: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 7281: 17.3356, expected 26.1401
2024-02-16 16:05:56.905669: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 31490: 46.9041, expected 35.7881
2024-02-16 16:05:56.905718: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 42098: 92.9747, expected 82.6593
2024-02-16 16:05:56.905856: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 77426: 14.8142, expected 10.7041
2024-02-16 16:05:56.905945: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 99275: 61.8406, expected 51.012
2024-02-16 16:05:56.905964: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 101673: 54.7939, expected 64.0151
2024-02-16 16:05:56.906010: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 111729: -23.2614, expected -28.2862
2024-02-16 16:05:56.906029: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 114033: -54.0424, expected -44.5037
2024-02-16 16:05:56.906077: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 124619: -35.2993, expected -46.6975
2024-02-16 16:05:56.906086: E external/xla/xla/service/gpu/triton_autotuner.cc:816] Results do not match the reference. This is likely a bug/unexpected loss of precision.
2024-02-16 16:05:57.417363: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 3557: -12.5601, expected -8.76762
2024-02-16 16:05:57.417414: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 7281: 17.2671, expected 26.1401
2024-02-16 16:05:57.417517: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 31490: 46.9089, expected 35.7881
2024-02-16 16:05:57.417572: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 42098: 92.8985, expected 82.6593
2024-02-16 16:05:57.417726: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 77426: 14.7896, expected 10.7041
2024-02-16 16:05:57.417815: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 99275: 61.807, expected 51.012
2024-02-16 16:05:57.417834: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 101673: 54.8168, expected 64.0151
2024-02-16 16:05:57.417882: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 111729: -23.3462, expected -28.2862
2024-02-16 16:05:57.417901: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 114033: -54.0808, expected -44.5037
2024-02-16 16:05:57.417951: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 124619: -35.2514, expected -46.6975
2024-02-16 16:05:57.417961: E external/xla/xla/service/gpu/triton_autotuner.cc:816] Results do not match the reference. This is likely a bug/unexpected loss of precision.
2024-02-16 16:05:57.460488: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 3557: -12.5601, expected -8.76762
2024-02-16 16:05:57.460531: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 7281: 17.2671, expected 26.1401
2024-02-16 16:05:57.460638: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 31490: 46.9089, expected 35.7881
2024-02-16 16:05:57.460692: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 42098: 92.8985, expected 82.6593
2024-02-16 16:05:57.460830: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 77426: 14.7896, expected 10.7041
2024-02-16 16:05:57.460919: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 99275: 61.807, expected 51.012
2024-02-16 16:05:57.460938: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 101673: 54.8168, expected 64.0151
2024-02-16 16:05:57.460985: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 111729: -23.3462, expected -28.2862
2024-02-16 16:05:57.461003: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 114033: -54.0808, expected -44.5037
2024-02-16 16:05:57.461052: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 124619: -35.2514, expected -46.6975
2024-02-16 16:05:57.461062: E external/xla/xla/service/gpu/triton_autotuner.cc:816] Results do not match the reference. This is likely a bug/unexpected loss of precision.
2024-02-16 16:05:57.520325: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 3557: -12.5601, expected -8.76762
2024-02-16 16:05:57.520366: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 7281: 17.2671, expected 26.1401
2024-02-16 16:05:57.520472: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 31490: 46.9089, expected 35.7881
2024-02-16 16:05:57.520523: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 42098: 92.8985, expected 82.6593
2024-02-16 16:05:57.520675: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 77426: 14.7896, expected 10.7041
2024-02-16 16:05:57.520765: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 99275: 61.807, expected 51.012
2024-02-16 16:05:57.520784: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 101673: 54.8168, expected 64.0151
2024-02-16 16:05:57.520830: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 111729: -23.3462, expected -28.2862
2024-02-16 16:05:57.520849: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 114033: -54.0808, expected -44.5037
2024-02-16 16:05:57.520897: E external/xla/xla/service/gpu/buffer_comparator.cc:149] Difference at 124619: -35.2514, expected -46.6975
2024-02-16 16:05:57.520907: E external/xla/xla/service/gpu/triton_autotuner.cc:816] Results do not match the reference. This is likely a bug/unexpected loss of precision.
The code for the Neural Network layer that seems to be at fault more or less (I'm giving a simplified version without inheritance etc.) comes down to
import jax
jax.config.update("jax_default_device", jax.devices('cpu')[0])
from jax import numpy as jnp
import equinox as eqx
class SirenLayer(eqx.Module):
weights: jax.Array
biases: jax.Array
activation_kwargs: dict = eqx.field(static=True)
def __call__(self, x):
wxb = (self.weights@x + self.biases, )
return self.activation_function(*wxb)
def activation_function(self, *args):
return self._activation_function(*args, **self.activation_kwargs)
@staticmethod
def _activation_function(x, w0):
return jnp.sin(w0*x)
@classmethod
def from_config(cls, in_size:int, out_size:int, *, key:jax.Array, is_first_layer:bool, w0:float):
w_key, b_key = jax.random.split(key)
if is_first_layer:
lim = 1./in_size# from https://github.com/vsitzmann/siren/blob/4df34baee3f0f9c8f351630992c1fe1f69114b5f/modules.py#L630
else:
lim = jnp.sqrt(6./in_size)/w0 # from https://arxiv.org/pdf/2006.09661.pdf subsection.3.2 and appendix 1.5 and https://github.com/vsitzmann/siren/blob/4df34baee3f0f9c8f351630992c1fe1f69114b5f/modules.py#L627
weight = jax.random.uniform(
key=w_key,
shape=(out_size, in_size),
minval=-lim,
maxval=lim
)
bias = jax.random.uniform(
key=b_key,
shape=(out_size,),
minval=-1,
maxval=1
)
bias_factor = jnp.pi/jnp.sqrt(jnp.sum(jnp.square(weight), axis=1)) # from https://arxiv.org/pdf/2102.02611.pdf page 6 third paragaph
bias = bias_factor * bias
return cls(weight, bias, {'w0': w0})
One of the networks that would sometimes give this error had an input size of 3, a linear output layer with output size 1, a hidden size of 384, and 6 layers total (including final linear layer and input layer). The value of w0 used was 16.
System info (python version, jaxlib version, accelerator, etc.)
jax: 0.4.24
jaxlib: 0.4.24
numpy: 1.26.4
python: 3.10.13 (main, Sep 11 2023, 13:44:35) [GCC 11.2.0]
jax.devices (1 total, 1 local): [cuda(id=0)]
process_count: 1
$ nvidia-smi
Tue Mar 5 16:06:24 2024
+---------------------------------------------------------------------------------------+
| NVIDIA-SMI 545.23.08 Driver Version: 545.23.08 CUDA Version: 12.3 |
|-----------------------------------------+----------------------+----------------------+
| GPU Name Persistence-M | Bus-Id Disp.A | Volatile Uncorr. ECC |
| Fan Temp Perf Pwr:Usage/Cap | Memory-Usage | GPU-Util Compute M. |
| | | MIG M. |
|=========================================+======================+======================|
| 0 NVIDIA A100-SXM4-40GB On | 00000000:31:00.0 Off | Off |
| N/A 31C P0 72W / 400W | 31124MiB / 40960MiB | 0% Default |
| | | Disabled |
+-----------------------------------------+----------------------+----------------------+
+---------------------------------------------------------------------------------------+
| Processes: |
| GPU GI CI PID Type Process name GPU Memory |
| ID ID Usage |
|=======================================================================================|
| 0 N/A N/A 2797663 C python 31110MiB |
+---------------------------------------------------------------------------------------+
The version of Equinox I use (in case it's important) is 0.11.3.
I'm also seeing this error when using sine activations. However, it appears to only occur when jitting the function.
I'm running into the same issue when using sine activation function. Just wanted to check if you guys were able to find any solution to this? Thanks!
I used this repo instead of my network and no longer had the issue. Perhaps something to do with how the network was defined.
The problem disappeared for me after updating JAX at some point. No idea what changed, but I'm happy my code works 😅
Op do 5 sep 2024 om 18:42 schreef Mukesh Ghimire @.***>:
I'm running into the same issue when using sine activation function. Just wanted to check if you guys were able to find any solution to this? Thanks!
— Reply to this email directly, view it on GitHub https://github.com/google/jax/issues/20075#issuecomment-2332194573, or unsubscribe https://github.com/notifications/unsubscribe-auth/AMWCTIGOWZX7MJ5BWHGR4YLZVCCXHAVCNFSM6AAAAABEHMJV7GVHI2DSMVQWIX3LMV43OSLTON2WKQ3PNVWWK3TUHMZDGMZSGE4TINJXGM . You are receiving this because you authored the thread.Message ID: @.***>
Just curious if the models performed differently when you had the issue vs now? Appreciate you responding.
I'm still getting the error when the weight $w_0 = 30$ in the activation.