Spaces:
Sleeping
Sleeping
File size: 6,426 Bytes
9bdaa77 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 |
# Copyright 2022 DeepMind Technologies Limited. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""MLP to compute basic linear functions of one-hot encoded integers."""
from typing import Callable
import numpy as np
from tracr.craft import bases
from tracr.craft import transformers
from tracr.craft import vectorspace_fns
_ONE_SPACE = bases.VectorSpaceWithBasis.from_names(["one"])
def map_categorical_mlp(
input_space: bases.VectorSpaceWithBasis,
output_space: bases.VectorSpaceWithBasis,
operation: Callable[[bases.BasisDirection], bases.BasisDirection],
) -> transformers.MLP:
"""Returns an MLP that encodes any categorical function of a single variable f(x).
The hidden layer is the identity and output combines this with a lookup table
output_k = sum(f(i)*input_i for all i in input space)
Args:
input_space: space containing the input x.
output_space: space containing possible outputs.
operation: A function operating on basis directions.
"""
def operation_fn(direction):
if direction in input_space:
output_direction = operation(direction)
if output_direction in output_space:
return output_space.vector_from_basis_direction(output_direction)
return output_space.null_vector()
first_layer = vectorspace_fns.Linear.from_action(input_space, output_space,
operation_fn)
second_layer = vectorspace_fns.project(output_space, output_space)
return transformers.MLP(first_layer, second_layer)
def map_categorical_to_numerical_mlp(
input_space: bases.VectorSpaceWithBasis,
output_space: bases.VectorSpaceWithBasis,
operation: Callable[[bases.Value], float],
) -> transformers.MLP:
"""Returns an MLP to compute f(x) from a categorical to a numerical variable.
The hidden layer is the identity and output combines this with a lookup table
output = sum(f(i)*input_i for all i in input space)
Args:
input_space: Vector space containing the input x.
output_space: Vector space to write the numerical output to.
operation: A function operating on basis directions.
"""
bases.ensure_dims(output_space, num_dims=1, name="output_space")
out_vec = output_space.vector_from_basis_direction(output_space.basis[0])
def operation_fn(direction):
if direction in input_space:
return operation(direction.value) * out_vec
return output_space.null_vector()
first_layer = vectorspace_fns.Linear.from_action(input_space, output_space,
operation_fn)
second_layer = vectorspace_fns.project(output_space, output_space)
return transformers.MLP(first_layer, second_layer)
def sequence_map_categorical_mlp(
input1_space: bases.VectorSpaceWithBasis,
input2_space: bases.VectorSpaceWithBasis,
output_space: bases.VectorSpaceWithBasis,
operation: Callable[[bases.BasisDirection, bases.BasisDirection],
bases.BasisDirection],
one_space: bases.VectorSpaceWithBasis = _ONE_SPACE,
hidden_name: bases.Name = "__hidden__",
) -> transformers.MLP:
"""Returns an MLP that encodes a categorical function of two variables f(x, y).
The hidden layer of the MLP computes the logical and of all input directions
hidden_i_j = ReLU(x_i+x_j-1)
And the output combines this with a lookup table
output_k = sum(f(i, j)*hidden_i_j for all i,j in input space)
Args:
input1_space: Vector space containing the input x.
input2_space: Vector space containing the input y.
output_space: Vector space to write outputs to.
operation: A function operating on basis directions.
one_space: a reserved 1-d space that always contains a 1.
hidden_name: Name for hidden dimensions.
"""
bases.ensure_dims(one_space, num_dims=1, name="one_space")
if not set(input1_space.basis).isdisjoint(input2_space.basis):
raise ValueError("Input spaces to a SequenceMap must be disjoint. "
"If input spaces are the same, use Map instead!")
input_space = bases.direct_sum(input1_space, input2_space, one_space)
def to_hidden(x, y):
return bases.BasisDirection(hidden_name, (x.name, x.value, y.name, y.value))
def from_hidden(h):
x_name, x_value, y_name, y_value = h.value
x_dir = bases.BasisDirection(x_name, x_value)
y_dir = bases.BasisDirection(y_name, y_value)
return x_dir, y_dir
hidden_dir = []
for dir1 in input1_space.basis:
for dir2 in input2_space.basis:
hidden_dir.append(to_hidden(dir1, dir2))
hidden_space = bases.VectorSpaceWithBasis(hidden_dir)
def logical_and(direction):
if direction in one_space:
out = bases.VectorInBasis(hidden_space.basis,
-np.ones(hidden_space.num_dims))
elif direction in input1_space:
dir1 = direction
out = hidden_space.null_vector()
for dir2 in input2_space.basis:
out += hidden_space.vector_from_basis_direction(to_hidden(dir1, dir2))
else:
dir2 = direction
out = hidden_space.null_vector()
for dir1 in input1_space.basis:
out += hidden_space.vector_from_basis_direction(to_hidden(dir1, dir2))
return out
first_layer = vectorspace_fns.Linear.from_action(input_space, hidden_space,
logical_and)
def operation_fn(direction):
dir1, dir2 = from_hidden(direction)
output_direction = operation(dir1, dir2)
if output_direction in output_space:
return output_space.vector_from_basis_direction(output_direction)
else:
return output_space.null_vector()
second_layer = vectorspace_fns.Linear.from_action(hidden_space, output_space,
operation_fn)
return transformers.MLP(first_layer, second_layer)
|