Optimizing algorithms for hardware devices

A method performed by one or more computers for obtaining an optimized algorithm that (i) is functionally equivalent to a target algorithm and (ii) optimizes one or more target properties when executed on a target set of one or more hardware devices, the method comprising:

• initializing a target tensor representing the target algorithm;
• generating, using a neural network having a plurality of network parameters, a tensor decomposition of the target tensor that parametrizes a candidate algorithm, wherein the neural network is configured to receive a state of a tensor as input and process the input in accordance with the network parameters to generate a network output comprising a policy for applying modifications to the tensor, and wherein generating the tensor decomposition comprises:
• for each step in a sequence of steps:
• obtaining a current state of the target tensor;
• determining a target network output for the current state by performing a tree search of a state tree having nodes representing states of the target tensor starting from a root node that represents the current state, wherein the tree search is guided by the neural network in accordance with the network parameters;
• applying a modification to the current state of the target tensor using the target network output for the current state;
• and
• determining whether to terminate the sequence based, at least in part, on whether the target tensor equals a zero tensor after the modification is applied;
• and
• generating the tensor decomposition from the modifications applied to the target tensor at each step in the sequence of steps, wherein the tensor decomposition is a factorization of the target tensor into a linear combination of rank-one terms that each represents a single linearly independent direction in a tensor product space and has a rank of one, wherein the rank-one terms of the factorization of the target tensor correspond to the modifications applied to the target tensor at each step in the sequence of steps;
• generating target property values for each of the target properties when executing the candidate algorithm on the target set of hardware devices;
• determining a benchmarking score for the tensor decomposition based on the target property values of the candidate algorithm;
• generating a training example from the tensor decomposition and the benchmarking score, the training example including the state of the tensor and the target network output at each step in the sequence of steps;
• and
• storing, in a training data store, the training example for use in updating the network parameters of the neural network.