MPS implementation in MacOS Sonoma 14.1.2+ raises warning errors on tensor anec.gain_offset_control shape

[cregouby]

Current situation

Right after MacOS upgrade to Sonoma 14.1.2, the MPS implementation runs with log on tabnet_pretrain() and tabnet_fit() with the following log

loc("outputTensor"("(mpsFileLoc): /AppleInternal/Library/BuildRoots/495c257e-668e-11ee-93ce-926038f30c31/Library/Caches/com.apple.xbs/Sources/MetalPerformanceShadersGraph/mpsgraph/MetalPerformanceShadersGraph/Core/Files/MPSGraphUtilities.mm":745:0)): error: 'anec.gain_offset_control' op result #0 must be 4D/5D memref of 16-bit float or 8-bit signed integer or 8-bit unsigned integer values, but got 'memref<1x256x1x73xi1>'

It seems it can be treated as a warning (end-user perspective)
Please report any occurrence where the training fails.

ReprEx

data("ames", package = "modeldata")
ids <- sample(nrow(ames), 256)
small_ames <- ames[ids,]
x <- ames[ids,-which(names(ames) == "Sale_Price")]
y <- ames[ids,]$Sale_Price

# ames common models
library(tabnet)
ames_pretrain <- tabnet_pretrain(x, y, epoch = 2, checkpoint_epochs = 1,
                                 verbose = FALSE)
#> loc("outputTensor"("(mpsFileLoc): /AppleInternal/Library/BuildRoots/495c257e-668e-11ee-93ce-926038f30c31/Library/Caches/com.apple.xbs/Sources/Metal
#> PerformanceShadersGraph/mpsgraph/MetalPerformanceShadersGraph/Core/Files/MPSGraphUtilities.mm":745:0)): error: 'anec.gain_offset_control' op result #0
#> must be 4D/5D memref of 16-bit float or 8-bit signed integer or 8-bit unsigned integer values, but got 'memref<1x256x1x73xi1>'

^{Created on 2024-01-05 with reprex v2.0.2}