keras/src/layers/preprocessing/index_lookup.py/IndexLookup/_expand_dims
class IndexLookup:    def _expand_dims(self, inputs, axis):
        if isinstance(inputs, tf.SparseTensor):
            return tf.sparse.expand_dims(inputs, axis)
        else:
            return tf.expand_dims(inputs, axis)