Mathy's embeddings model takes in a window of observations and outputs a sequence of the same length with fixed-size learned embeddings for each token in the sequence.
The Mathy embeddings model is a stateful model that predicts over sequences. This complicates the process of collecting observations to feed to the model, but allows richer input features than would be available from the simpler state representation.
The model accepts an encoded sequence of tokens and values extracted from the current state's expression tree.
Observations to Embeddings¶
You can instantiate a model and produce untrained embeddings:
from mathy import envs from mathy.agents.base_config import BaseConfig from mathy.agents.embedding import MathyEmbedding from mathy.env import MathyEnv from mathy.state import MathyObservation, observations_to_window args = BaseConfig() env: MathyEnv = envs.PolySimplify() observation: MathyObservation = env.state_to_observation(env.get_initial_state()) model = MathyEmbedding(args) # output shape is: [num_observations, max_nodes_len, embedding_dimensions] inputs = observations_to_window([observation, observation]).to_inputs() embeddings, attentions = model(inputs) # We provided two observations in a sequence assert embeddings.shape == 2 # There are as many outputs as input sequences assert embeddings.shape == len(observation.nodes) # Outputs vectors with the provided embedding units assert embeddings.shape[-1] == args.embedding_units # The attention output is a grid of [num_observations, max_nodes_len, max_nodes_len] assert attentions.shape == [2, len(observation.nodes), len(observation.nodes)]