"""Provides Geometric Algebra Keras layers."""
from typing import List, Union
import tensorflow as tf
from tensorflow.keras import layers
from tensorflow.keras import (
initializers, activations, regularizers, constraints
)
from tensorflow.keras.utils import register_keras_serializable
from .blades import BladeKind
from .tfga import GeometricAlgebra
[docs]class GeometricAlgebraLayer(layers.Layer):
def __init__(self, algebra: GeometricAlgebra, **kwargs):
self.algebra = algebra
super().__init__(**kwargs)
[docs] @classmethod
def from_config(cls, config):
# Create algebra if necessary (should only occur once, assumes that
# config is actually mutable).
if "algebra" not in config:
assert "metric" in config
config["algebra"] = GeometricAlgebra(config["metric"])
del config["metric"]
return cls(**config)
[docs] def get_config(self):
# Store metric of the algebra. In from_config() we will recreate the
# algebra from the metric.
config = super().get_config()
config.update({
"metric": self.algebra.metric.numpy()
})
return config
[docs]@register_keras_serializable(package="TFGA")
class TensorToGeometric(GeometricAlgebraLayer):
"""Layer for converting tensors with given blade indices to
geometric algebra tensors.
Args:
algebra: GeometricAlgebra instance to use
blade_indices: blade indices to interpret the last axis of the
input tensor as
"""
def __init__(self, algebra: GeometricAlgebra, blade_indices: List[int],
**kwargs):
super().__init__(algebra=algebra, **kwargs)
self.blade_indices = tf.convert_to_tensor(
blade_indices, dtype=tf.int64)
[docs] def compute_output_shape(self, input_shape):
return tf.TensorShape([*input_shape[:-1], self.algebra.num_blades])
[docs] def call(self, inputs):
return self.algebra.from_tensor(inputs, blade_indices=self.blade_indices)
[docs] def get_config(self):
config = super().get_config()
config.update({
"blade_indices": self.blade_indices.numpy()
})
return config
[docs]@register_keras_serializable(package="TFGA")
class TensorWithKindToGeometric(GeometricAlgebraLayer):
"""Layer for converting tensors with given blade kind to
geometric algebra tensors.
Args:
algebra: GeometricAlgebra instance to use
kind: blade kind indices to interpret the last axis of the
input tensor as
"""
def __init__(self, algebra: GeometricAlgebra, kind: BladeKind,
**kwargs):
super().__init__(algebra=algebra, **kwargs)
self.kind = kind
[docs] def compute_output_shape(self, input_shape):
return tf.TensorShape([*input_shape[:-1], self.algebra.get_kind_blade_indices(self.kind).shape[0]])
[docs] def call(self, inputs):
return self.algebra.from_tensor_with_kind(inputs, kind=self.kind)
[docs] def get_config(self):
config = super().get_config()
config.update({
"kind": self.kind
})
return config
[docs]@register_keras_serializable(package="TFGA")
class GeometricToTensor(GeometricAlgebraLayer):
"""Layer for extracting given blades from geometric algebra tensors.
Args:
algebra: GeometricAlgebra instance to use
blade_indices: blade indices to extract
"""
def __init__(self, algebra: GeometricAlgebra, blade_indices: List[int],
**kwargs):
super().__init__(algebra=algebra, **kwargs)
self.blade_indices = tf.convert_to_tensor(
blade_indices, dtype=tf.int64)
[docs] def compute_output_shape(self, input_shape):
return tf.TensorShape([*input_shape[:-1], self.blade_indices.shape[0]])
[docs] def call(self, inputs):
return tf.gather(inputs, self.blade_indices, axis=-1)
[docs] def get_config(self):
config = super().get_config()
config.update({
"blade_indices": self.blade_indices.numpy()
})
return config
[docs]@register_keras_serializable(package="TFGA")
class GeometricToTensorWithKind(GeometricToTensor):
"""Layer for extracting blades of a kind from geometric algebra tensors.
Args:
algebra: GeometricAlgebra instance to use
kind: blade indices of kind to extract
"""
def __init__(self, algebra: GeometricAlgebra, kind: BladeKind,
**kwargs):
blade_indices = algebra.get_kind_blade_indices(kind)
super().__init__(algebra=algebra, blade_indices=blade_indices,
**kwargs)
[docs]@register_keras_serializable(package="TFGA")
class GeometricProductDense(GeometricAlgebraLayer):
"""Analagous to Keras' Dense layer but using multivector-valued matrices
instead of scalar ones and geometric multiplication instead of standard
multiplication.
Args:
algebra: GeometricAlgebra instance to use for the parameters
blade_indices_kernel: Blade indices to use for the kernel parameter
blade_indices_bias: Blade indices to use for the bias parameter (if used)
"""
def __init__(
self,
algebra: GeometricAlgebra,
units: int,
blade_indices_kernel: List[int],
blade_indices_bias: Union[None, List[int]] = None,
activation=None,
use_bias=True,
kernel_initializer="glorot_uniform",
bias_initializer="zeros",
kernel_regularizer=None,
bias_regularizer=None,
activity_regularizer=None,
kernel_constraint=None,
bias_constraint=None,
**kwargs
):
super().__init__(algebra=algebra, activity_regularizer=activity_regularizer, **kwargs)
self.units = units
self.blade_indices_kernel = tf.convert_to_tensor(
blade_indices_kernel, dtype_hint=tf.int64)
if use_bias:
self.blade_indices_bias = tf.convert_to_tensor(
blade_indices_bias, dtype_hint=tf.int64)
self.activation = activations.get(activation)
self.use_bias = use_bias
self.kernel_initializer = initializers.get(kernel_initializer)
self.bias_initializer = initializers.get(bias_initializer)
self.kernel_regularizer = regularizers.get(kernel_regularizer)
self.bias_regularizer = regularizers.get(bias_regularizer)
self.kernel_constraint = constraints.get(kernel_constraint)
self.bias_constraint = constraints.get(bias_constraint)
[docs] def build(self, input_shape: tf.TensorShape):
self.num_input_units = input_shape[-2]
shape_kernel = [
self.units,
self.num_input_units,
self.blade_indices_kernel.shape[0]
]
self.kernel = self.add_weight(
"kernel",
shape=shape_kernel,
initializer=self.kernel_initializer,
regularizer=self.kernel_regularizer,
constraint=self.kernel_constraint,
dtype=self.dtype,
trainable=True
)
if self.use_bias:
shape_bias = [self.units, self.blade_indices_bias.shape[0]]
self.bias = self.add_weight(
"bias",
shape=shape_bias,
initializer=self.bias_initializer,
regularizer=self.bias_regularizer,
constraint=self.bias_constraint,
dtype=self.dtype,
trainable=True
)
else:
self.bias = None
self.built = True
[docs] def compute_output_shape(self, input_shape):
return tf.TensorShape([*input_shape[:-2], self.units, self.algebra.num_blades])
[docs] def call(self, inputs):
w_geom = self.algebra.from_tensor(
self.kernel, self.blade_indices_kernel)
# Perform a matrix-multiply, but using geometric product instead of
# standard multiplication. To do this we do the geometric product
# elementwise and then sum over the common axis.
# [..., 1, I, X] * [..., O, I, X] -> [..., O, I, X] -> [..., O, X]
inputs_expanded = tf.expand_dims(inputs, axis=inputs.shape.ndims - 2)
result = tf.reduce_sum(self.algebra.geom_prod(
inputs_expanded, w_geom), axis=-2)
if self.bias is not None:
b_geom = self.algebra.from_tensor(
self.bias, self.blade_indices_bias)
result += b_geom
return self.activation(result)
[docs] def get_config(self):
config = super().get_config()
config.update({
"blade_indices_kernel":
self.blade_indices_kernel.numpy(),
"blade_indices_bias":
self.blade_indices_bias.numpy(),
"units":
self.units,
"activation":
activations.serialize(self.activation),
"use_bias":
self.use_bias,
"kernel_initializer":
initializers.serialize(self.kernel_initializer),
"bias_initializer":
initializers.serialize(self.bias_initializer),
"kernel_regularizer":
regularizers.serialize(self.kernel_regularizer),
"bias_regularizer":
regularizers.serialize(self.bias_regularizer),
"activity_regularizer":
regularizers.serialize(self.activity_regularizer),
"kernel_constraint":
constraints.serialize(self.kernel_constraint),
"bias_constraint":
constraints.serialize(self.bias_constraint)
})
return config
[docs]@register_keras_serializable(package="TFGA")
class GeometricSandwichProductDense(GeometricProductDense):
"""Analagous to Keras' Dense layer but using multivector-valued matrices
instead of scalar ones and geometric sandwich multiplication instead of
standard multiplication.
Args:
algebra: GeometricAlgebra instance to use for the parameters
blade_indices_kernel: Blade indices to use for the kernel parameter
blade_indices_bias: Blade indices to use for the bias parameter (if used)
"""
def __init__(
self, algebra, units, blade_indices_kernel, blade_indices_bias=None,
activation=None, use_bias=True, kernel_initializer="glorot_uniform",
bias_initializer="zeros", kernel_regularizer=None,
bias_regularizer=None, activity_regularizer=None,
kernel_constraint=None, bias_constraint=None, **kwargs
):
super().__init__(
algebra, units,
blade_indices_kernel,
blade_indices_bias=blade_indices_bias,
activation=activation,
use_bias=use_bias,
kernel_initializer=kernel_initializer,
bias_initializer=bias_initializer,
kernel_regularizer=kernel_regularizer,
bias_regularizer=bias_regularizer,
activity_regularizer=activity_regularizer,
kernel_constraint=kernel_constraint,
bias_constraint=bias_constraint, **kwargs
)
[docs] def call(self, inputs):
w_geom = self.algebra.from_tensor(
self.kernel, self.blade_indices_kernel)
# Same as GeometricProductDense but using R*x*~R instead of just R*x
inputs_expanded = tf.expand_dims(inputs, axis=inputs.shape.ndims - 2)
result = tf.reduce_sum(
self.algebra.geom_prod(
w_geom,
self.algebra.geom_prod(
inputs_expanded,
self.algebra.reversion(w_geom)
)
),
axis=-2
)
if self.bias is not None:
b_geom = self.algebra.from_tensor(
self.bias, self.blade_indices_bias)
result += b_geom
return self.activation(result)
[docs]@register_keras_serializable(package="TFGA")
class GeometricProductElementwise(GeometricAlgebraLayer):
"""Performs the elementwise geometric product with a list of multivectors
with as many elements as there are input units.
Args:
algebra: GeometricAlgebra instance to use for the parameters
blade_indices_kernel: Blade indices to use for the kernel parameter
blade_indices_bias: Blade indices to use for the bias parameter (if used)
"""
def __init__(
self,
algebra: GeometricAlgebra,
blade_indices_kernel: List[int],
blade_indices_bias: Union[None, List[int]] = None,
activation=None,
use_bias=True,
kernel_initializer="glorot_uniform",
bias_initializer="zeros",
kernel_regularizer=None,
bias_regularizer=None,
activity_regularizer=None,
kernel_constraint=None,
bias_constraint=None,
**kwargs
):
super().__init__(algebra=algebra, activity_regularizer=activity_regularizer, **kwargs)
self.blade_indices_kernel = tf.convert_to_tensor(
blade_indices_kernel, dtype_hint=tf.int64)
if use_bias:
self.blade_indices_bias = tf.convert_to_tensor(
blade_indices_bias, dtype_hint=tf.int64)
self.activation = activations.get(activation)
self.use_bias = use_bias
self.kernel_initializer = initializers.get(kernel_initializer)
self.bias_initializer = initializers.get(bias_initializer)
self.kernel_regularizer = regularizers.get(kernel_regularizer)
self.bias_regularizer = regularizers.get(bias_regularizer)
self.kernel_constraint = constraints.get(kernel_constraint)
self.bias_constraint = constraints.get(bias_constraint)
[docs] def build(self, input_shape: tf.TensorShape):
self.num_input_units = input_shape[-2]
shape_kernel = [
self.num_input_units,
self.blade_indices_kernel.shape[0]
]
self.kernel = self.add_weight(
"kernel",
shape=shape_kernel,
initializer=self.kernel_initializer,
regularizer=self.kernel_regularizer,
constraint=self.kernel_constraint,
dtype=self.dtype,
trainable=True
)
if self.use_bias:
shape_bias = [self.num_input_units,
self.blade_indices_bias.shape[0]]
self.bias = self.add_weight(
"bias",
shape=shape_bias,
initializer=self.bias_initializer,
regularizer=self.bias_regularizer,
constraint=self.bias_constraint,
dtype=self.dtype,
trainable=True
)
else:
self.bias = None
self.built = True
[docs] def compute_output_shape(self, input_shape):
return tf.TensorShape([*input_shape[:-1], self.algebra.num_blades])
[docs] def call(self, inputs):
w_geom = self.algebra.from_tensor(
self.kernel, self.blade_indices_kernel)
# Elementwise multiplication for each unit with a multivector.
# [..., U, X] * [U, X] -> [..., U, X]
result = self.algebra.geom_prod(inputs, w_geom)
if self.bias is not None:
b_geom = self.algebra.from_tensor(
self.bias, self.blade_indices_bias)
result += b_geom
return self.activation(result)
[docs] def get_config(self):
config = super().get_config()
config.update({
"blade_indices_kernel":
self.blade_indices_kernel.numpy(),
"blade_indices_bias":
self.blade_indices_bias.numpy(),
"activation":
activations.serialize(self.activation),
"use_bias":
self.use_bias,
"kernel_initializer":
initializers.serialize(self.kernel_initializer),
"bias_initializer":
initializers.serialize(self.bias_initializer),
"kernel_regularizer":
regularizers.serialize(self.kernel_regularizer),
"bias_regularizer":
regularizers.serialize(self.bias_regularizer),
"activity_regularizer":
regularizers.serialize(self.activity_regularizer),
"kernel_constraint":
constraints.serialize(self.kernel_constraint),
"bias_constraint":
constraints.serialize(self.bias_constraint)
})
return config
[docs]@register_keras_serializable(package="TFGA")
class GeometricSandwichProductElementwise(GeometricProductElementwise):
"""Performs the elementwise geometric sandwich product with a list of
multivectors with as many elements as there are input units.
Args:
algebra: GeometricAlgebra instance to use for the parameters
blade_indices_kernel: Blade indices to use for the kernel parameter
blade_indices_bias: Blade indices to use for the bias parameter (if used)
"""
def __init__(
self, algebra, blade_indices_kernel, blade_indices_bias=None,
activation=None, use_bias=True, kernel_initializer="glorot_uniform",
bias_initializer="zeros", kernel_regularizer=None,
bias_regularizer=None, activity_regularizer=None,
kernel_constraint=None, bias_constraint=None, **kwargs
):
super().__init__(
algebra,
blade_indices_kernel,
blade_indices_bias=blade_indices_bias,
activation=activation,
use_bias=use_bias,
kernel_initializer=kernel_initializer,
bias_initializer=bias_initializer,
kernel_regularizer=kernel_regularizer,
bias_regularizer=bias_regularizer,
activity_regularizer=activity_regularizer,
kernel_constraint=kernel_constraint,
bias_constraint=bias_constraint, **kwargs
)
[docs] def call(self, inputs):
w_geom = self.algebra.from_tensor(
self.kernel, self.blade_indices_kernel)
# Elementwise multiplication Rx~R for each unit with a multivector.
# [..., U, X] * [U, X] -> [..., U, X]
result = self.algebra.geom_prod(
w_geom,
self.algebra.geom_prod(
inputs,
self.algebra.reversion(w_geom)
)
)
if self.bias is not None:
b_geom = self.algebra.from_tensor(
self.bias, self.blade_indices_bias)
result += b_geom
return self.activation(result)
[docs]@register_keras_serializable(package="TFGA")
class GeometricProductConv1D(GeometricAlgebraLayer):
"""Analagous to Keras' Conv1D layer but using multivector-valued kernels
instead of scalar ones and geometric product instead of
standard multiplication.
Args:
algebra: GeometricAlgebra instance to use for the parameters
filters: How many channels the output will have
kernel_size: Size for the convolution kernel
stride: Stride to use for the convolution
padding: "SAME" (zero-pad input length so output
length == input length / stride) or "VALID" (no padding)
blade_indices_kernel: Blade indices to use for the kernel parameter
blade_indices_bias: Blade indices to use for the bias parameter (if used)
"""
def __init__(
self,
algebra: GeometricAlgebra,
filters: int,
kernel_size: int,
stride: int,
padding: str,
blade_indices_kernel: List[int],
blade_indices_bias: Union[None, List[int]] = None,
dilations: Union[None, int] = None,
activation=None,
use_bias=True,
kernel_initializer="glorot_uniform",
bias_initializer="zeros",
kernel_regularizer=None,
bias_regularizer=None,
activity_regularizer=None,
kernel_constraint=None,
bias_constraint=None,
**kwargs
):
super().__init__(
algebra=algebra,
activity_regularizer=activity_regularizer,
**kwargs
)
self.filters = filters
self.kernel_size = kernel_size
self.stride = stride
self.padding = padding
self.dilations = dilations
self.blade_indices_kernel = tf.convert_to_tensor(
blade_indices_kernel, dtype_hint=tf.int64)
if use_bias:
self.blade_indices_bias = tf.convert_to_tensor(
blade_indices_bias, dtype_hint=tf.int64)
self.activation = activations.get(activation)
self.use_bias = use_bias
self.kernel_initializer = initializers.get(kernel_initializer)
self.bias_initializer = initializers.get(bias_initializer)
self.kernel_regularizer = regularizers.get(kernel_regularizer)
self.bias_regularizer = regularizers.get(bias_regularizer)
self.kernel_constraint = constraints.get(kernel_constraint)
self.bias_constraint = constraints.get(bias_constraint)
[docs] def build(self, input_shape: tf.TensorShape):
# I: [..., S, C, B]
self.num_input_filters = input_shape[-2]
# K: [K, IC, OC, B]
shape_kernel = [
self.kernel_size,
self.num_input_filters,
self.filters,
self.blade_indices_kernel.shape[0]
]
self.kernel = self.add_weight(
"kernel",
shape=shape_kernel,
initializer=self.kernel_initializer,
regularizer=self.kernel_regularizer,
constraint=self.kernel_constraint,
dtype=self.dtype,
trainable=True
)
if self.use_bias:
shape_bias = [self.filters, self.blade_indices_bias.shape[0]]
self.bias = self.add_weight(
"bias",
shape=shape_bias,
initializer=self.bias_initializer,
regularizer=self.bias_regularizer,
constraint=self.bias_constraint,
dtype=self.dtype,
trainable=True
)
else:
self.bias = None
self.built = True
[docs] def call(self, inputs):
k_geom = self.algebra.from_tensor(
self.kernel, self.blade_indices_kernel)
result = self.algebra.geom_conv1d(
inputs, k_geom,
stride=self.stride, padding=self.padding,
dilations=self.dilations
)
if self.bias is not None:
b_geom = self.algebra.from_tensor(
self.bias, self.blade_indices_bias)
result += b_geom
return self.activation(result)
[docs] def get_config(self):
config = super().get_config()
config.update({
"filters":
self.filters,
"kernel_size":
self.kernel_size,
"stride":
self.stride,
"padding":
self.padding,
"dilations":
self.dilations,
"blade_indices_kernel":
self.blade_indices_kernel.numpy(),
"blade_indices_bias":
self.blade_indices_bias.numpy(),
"activation":
activations.serialize(self.activation),
"use_bias":
self.use_bias,
"kernel_initializer":
initializers.serialize(self.kernel_initializer),
"bias_initializer":
initializers.serialize(self.bias_initializer),
"kernel_regularizer":
regularizers.serialize(self.kernel_regularizer),
"bias_regularizer":
regularizers.serialize(self.bias_regularizer),
"activity_regularizer":
regularizers.serialize(self.activity_regularizer),
"kernel_constraint":
constraints.serialize(self.kernel_constraint),
"bias_constraint":
constraints.serialize(self.bias_constraint)
})
return config
[docs]@register_keras_serializable(package="TFGA")
class GeometricAlgebraExp(GeometricAlgebraLayer):
"""Calculates the exponential function of the input. Input must square to
a scalar.
Args:
algebra: GeometricAlgebra instance to use
square_scalar_tolerance: Tolerance to use for the square scalar check
or None if the check should be skipped
"""
def __init__(
self,
algebra: GeometricAlgebra,
square_scalar_tolerance: Union[float, None] = 1e-4,
**kwargs
):
super().__init__(algebra=algebra, **kwargs)
self.square_scalar_tolerance = square_scalar_tolerance
[docs] def compute_output_shape(self, input_shape):
return tf.TensorShape([*input_shape[:-1], self.algebra.num_blades])
[docs] def call(self, inputs):
return self.algebra.exp(
inputs, square_scalar_tolerance=self.square_scalar_tolerance
)
[docs] def get_config(self):
config = super().get_config()
config.update({
"square_scalar_tolerance": self.square_scalar_tolerance
})
return config