I tried using the CustomAdam() function (above). I runs, but it doesn't appear to be using the GPU (according to the performance meter). I'm on an iMac 27" with an AMD Radeon Pro 5700XT `import tensorflow as tf mnist = tf.keras.datasets.mnist (x_train, y_train), (x_test, y_test) = mnist.load_data() x_train, x_test = x_train / 255.0, x_test / 255.0 model = tf.keras.models.Sequential([ tf.keras.layers.Flatten(input_shape=(28, 28)), tf.keras.layers.Dense(128, activation='relu'), tf.keras.layers.Dropout(0.2), tf.keras.layers.Dense(10) ]) predictions = model(x_train[:1]).numpy() tf.nn.softmax(predictions).numpy() loss_fn = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True) loss_fn(y_train[:1], predictions).numpy() model.compile(optimizer = CustomAdam(), loss=loss_fn) model.fit(x_train, y_train, epochs=10)` Here's output: `2022-03-26 12:29:25.682832: I tensorflow/core/platform/cpu_feature_guard.cc:151] This TensorFlow binary is optimized with oneAPI Deep Neural Network Library (oneDNN) to use the following CPU instructions in performance-critical operations: SSE4.2 AVX AVX2 FMA To enable them in other operations, rebuild TensorFlow with the appropriate compiler flags. 2022-03-26 12:29:25.683388: I tensorflow/core/common_runtime/pluggable_device/pluggable_device_factory.cc:305] Could not identify NUMA node of platform GPU ID 0, defaulting to 0. Your kernel may not have been built with NUMA support. 2022-03-26 12:29:25.683664: I tensorflow/core/common_runtime/pluggable_device/pluggable_device_factory.cc:271] Created TensorFlow device (/job:localhost/replica:0/task:0/device:GPU:0 with 0 MB memory) -> physical PluggableDevice (device: 0, name: METAL, pci bus id: ) Metal device set to: AMD Radeon Pro 5700 XT Epoch 1/10 145/1875 [=>............................] - ETA: 1:35 - loss: 0.8532 [ ]: ​ `

Tried: optimizer=tfa.optimizers.RectifiedAdam(). It's not utilizing GPU on iMac 27" with AMD Radeon Pro 5700 XT in this code example: `inputs = keras.Input(shape=(sequence_length, raw_data.shape[-1])) x = layers.LSTM(32, recurrent_dropout=0.25)(inputs) x = layers.Dropout(0.5)(x) outputs = layers.Dense(1)(x) model = keras.Model(inputs, outputs) callbacks = [ keras.callbacks.ModelCheckpoint("jena_lstm_dropout.keras", save_best_only=True) ] model.compile(optimizer=tfa.optimizers.RectifiedAdam(), loss="mse", metrics=["mae"]) history = model.fit(train_dataset, epochs=50, validation_data=val_dataset, callbacks=callbacks)`

Show us the code ...

Confirmed. It's an Intel based iMac 27" (2020) with an AMD Radeo Pro 5700 XT GPU running OS X 12.3.1 % ipython Python 3.8.5 ... % pip show tensorflow-macos WARNING: Ignoring invalid distribution -umpy (/Users/davidlaxer/tensorflow-metal/lib/python3.8/site-packages) Name: tensorflow-macos Version: 2.8.0 Summary: TensorFlow is an open source machine learning framework for everyone. Home-page: https://www.tensorflow.org/ Author: Google Inc. Author-email: packages@tensorflow.org License: Apache 2.0 Location: /Users/davidlaxer/tensorflow-metal/lib/python3.8/site-packages Requires: absl-py, astunparse, flatbuffers, gast, google-pasta, grpcio, h5py, keras, keras-preprocessing, libclang, numpy, opt-einsum, protobuf, setuptools, six, tensorboard, termcolor, tf-estimator-nightly, typing-extensions, wrapt Required-by: (tensorflow-metal) (base) davidlaxer@x86_64-apple-darwin13 top2vec % pip show tensorflow-metal WARNING: Ignoring invalid distribution -umpy (/Users/davidlaxer/tensorflow-metal/lib/python3.8/site-packages) Name: tensorflow-metal Version: 0.4.0 Summary: TensorFlow acceleration for Mac GPUs. Home-page: https://developer.apple.com/metal/tensorflow-plugin/ Author: Author-email: License: MIT License. Copyright © 2020-2021 Apple Inc. All rights reserved. Location: /Users/davidlaxer/tensorflow-metal/lib/python3.8/site-packages Requires: six, wheel Required-by: (tensorflow-metal) (base) davidlaxer@x86_64-apple-darwin13 top2vec %

https://github.com/pytorch/pytorch/issues/77867

Does the M2 chip have 64 bit registers in the GPU?