[
  {
    "path": ".gitignore",
    "content": "venv\n__pycache__\n*.pyo\n*.pyc\nbuild\ndist\n*.egg-info\n*.swp\n"
  },
  {
    "path": "LICENSE",
    "content": "\"License\" shall mean the terms and conditions for use, reproduction, and distribution as defined by the text below.\n \n\"You\" (or \"Your\") shall mean an individual or Legal Entity exercising permissions granted by this License.\n \n\"Legal Entity\" shall mean the union of the acting entity and all other entities that control, are controlled by, or are under common control with that entity. For the purposes of this definition, \"control\" means (i) the power, direct or indirect, to cause the direction or management of such entity, whether by contract or otherwise, or (ii) ownership of fifty percent (50%) or more of the outstanding shares, or (iii) beneficial ownership of such entity.\n \n\"Source\" form shall mean the preferred form for making modifications, including but not limited to software source code, documentation source, and configuration files.\n \n\"Object\" form shall mean any form resulting from mechanical transformation or translation of a Source form, including but not limited to compiled object code, generated documentation, and conversions to other media types.\n \n\"Work\" shall mean the work of authorship, whether in Source or Object form, made available under this License.\n \nThis License governs use of the accompanying Work, and your use of the Work constitutes acceptance of this License.\n \nYou may use this Work for any non-commercial purpose, subject to the restrictions in this License. Some purposes which can be non-commercial are teaching, academic research, and personal experimentation. You may also distribute this Work with books or other teaching materials, or publish the Work on websites, that are intended to teach the use of the Work.\n \nYou may not use or distribute this Work, or any derivative works, outputs, or results from the Work, in any form for commercial purposes. Non-exhaustive examples of commercial purposes would be running business operations, licensing, leasing, or selling the Work, or distributing the Work for use with commercial products.\n \nYou may modify this Work and distribute the modified Work for non-commercial purposes, however, you may not grant rights to the Work or derivative works that are broader than or in conflict with those provided by this License. For example, you may not distribute modifications of the Work under terms that would permit commercial use, or under terms that purport to require the Work or derivative works to be sublicensed to others.\n\nIn return, we require that you agree:\n\n1. Not to remove any copyright or other notices from the Work.\n \n2. That if you distribute the Work in Source or Object form, you will include a verbatim copy of this License.\n \n3. That if you distribute derivative works of the Work in Source form, you do so only under a license that includes all of the provisions of this License and is not in conflict with this License, and if you distribute derivative works of the Work solely in Object form you do so only under a license that complies with this License.\n \n4. That if you have modified the Work or created derivative works from the Work, and distribute such modifications or derivative works, you will cause the modified files to carry prominent notices so that recipients know that they are not receiving the original Work. Such notices must state: (i) that you have changed the Work; and (ii) the date of any changes.\n \n5. If you publicly use the Work or any output or result of the Work, you will provide a notice with such use that provides any person who uses, views, accesses, interacts with, or is otherwise exposed to the Work (i) with information of the nature of the Work, (ii) with a link to the Work, and (iii) a notice that the Work is available under this License.\n \n6. THAT THE WORK COMES \"AS IS\", WITH NO WARRANTIES. THIS MEANS NO EXPRESS, IMPLIED OR STATUTORY WARRANTY, INCLUDING WITHOUT LIMITATION, WARRANTIES OF MERCHANTABILITY OR FITNESS FOR A PARTICULAR PURPOSE OR ANY WARRANTY OF TITLE OR NON-INFRINGEMENT. ALSO, YOU MUST PASS THIS DISCLAIMER ON WHENEVER YOU DISTRIBUTE THE WORK OR DERIVATIVE WORKS.\n \n7. THAT NEITHER UBER TECHNOLOGIES, INC. NOR ANY OF ITS AFFILIATES, SUPPLIERS, SUCCESSORS, NOR ASSIGNS WILL BE LIABLE FOR ANY DAMAGES RELATED TO THE WORK OR THIS LICENSE, INCLUDING DIRECT, INDIRECT, SPECIAL, CONSEQUENTIAL OR INCIDENTAL DAMAGES, TO THE MAXIMUM EXTENT THE LAW PERMITS, NO MATTER WHAT LEGAL THEORY IT IS BASED ON. ALSO, YOU MUST PASS THIS LIMITATION OF LIABILITY ON WHENEVER YOU DISTRIBUTE THE WORK OR DERIVATIVE WORKS.\n \n8. That if you sue anyone over patents that you think may apply to the Work or anyone's use of the Work, your license to the Work ends automatically.\n \n9. That your rights under the License end automatically if you breach it in any way.\n \n10. Uber Technologies, Inc. reserves all rights not expressly granted to you in this License.\n\n"
  },
  {
    "path": "README.md",
    "content": "# EvoGrad\n\nEvoGrad is a lightweight tool for differentiating through expectation,\nbuilt on top of PyTorch.\n\nTools that enable fast and flexible experimentation democratize and accelerate machine learning research.\nHowever, one field that so far has not been greatly impacted by automatic differentiation tools is evolutionary computation\nThe reason is that most evolutionary algorithms are gradient-free:\nthey do not follow any explicit mathematical gradient (i.e., the mathematically optimal local direction of improvement), and instead proceed through a generate-and-test heuristic.\nIn other words, they create new variants, test them out, and keep the best.\n\nRecent and exciting research in evolutionary algorithms for deep reinforcement learning, however, has highlighted how a specific class of evolutionary algorithms can benefit from auto-differentiation.\nWork from OpenAI demonstrated that a form of Natural Evolution Strategies (NES) is massively scalable, and competitive with modern deep reinforcement learning algorithms.\n\nEvoGrad enables fast prototyping of NES-like algorithms.\nWe believe there are many interesting algorithms yet to be discovered in this vein, and we hope this library will help to catalyze progress in the machine learning community.\n\n## Examples\n### Natural Evolution Strategies\nAs a first example, we’ll implement the simplified NES algorithm of [Salimans et al. (2017)](https://openai.com/blog/evolution-strategies/) in EvoGrad.\nEvoGrad provides several probability distributions which may be used in the expectation function.\nWe will use a normal distribution because it is the most common choice in practice.\n\nLet’s consider the problem of finding a fitness peak in a simple 1-D search space.\n\nWe can define our population distribution over this search space to be initially centered at 1.0, with a fixed variance of 0.05, with the following Python code:\n\n```python\nmu = torch.tensor([1.0], requires_grad=True)\np = Normal(mu, 0.05)\n```\n\nNext, let’s define a simple fitness function that rewards individuals for approaching the location 5.0 in the search space:\n\n```python\ndef fitness(xs):\n\treturn -(x - 5.0) ** 2\n```\n\nEach generation of evolution in NES takes samples from the population distribution and evaluates the fitness of each of those individual samples. Here we sample and evaluate 100 individuals from the distribution:\n\n```python\nsample = p.sample(n=100)\nfitnesses = fitness(sample)\n```\n\nOptionally, we can apply a [whitening transformation](https://en.wikipedia.org/wiki/Whitening_transformation) to the fitnesses (a form of pre-processing that often increases NES performance), like this:\n\n```python\nfitnesses = (fitnesses - fitnesses.mean()) / fitnesses.std()\n```\n\nNow we can use these calculated fitness values to estimate the mean fitness over our population distribution:\n\n```python\nmean = expectation(fitnesses, sample, p=p)\n```\n\nAlthough we could have estimated the mean value directly with the snippet:\n`mean = fitnesses.mean()`, what we gain by instead using the EvoGrad `expectation` function is the ability to backpropagate through mean.\nWe can then use the resulting auto-differentiated gradients to optimize the center of the 1D Gaussian population distribution (mu) through gradient descent (here, to increase the expected fitness value of the population):\n\n```python\nmean.backward()\nwith torch.no_grad():\n\tmu += alpha * mu.grad\n\tmu.grad.zero_()\n```\n\n### Maximizing Variance\nAs a more sophisticated example, rather than maximizing the mean fitness, we can maximize the variance of behaviors in the population.\nWhile fitness is a measure of quality for a fixed task, in some situations we want to prepare for the unknown, and instead might want our population to contain a diversity of behaviors that can easily be adapted to solve a wide range of possible future tasks.\n\nTo do so, we need a quantification of behavior, which we can call a behavior characterization. Similarly to how you can evaluate an individual parameter vector drawn from the population distribution to establish its fitness (e.g. how far does this controller cause a robot to walk?), you could evaluate such a draw and return some quantification of its behavior (e.g., what position does a robot controlled by this neural network locomote to?).\n\nFor this example, let’s choose a simple but illustrative, 1D behavior characterization, namely, the product of two sine waves (one with much faster frequency than the other):\n\n```python\ndef behavior(x):\n\treturn 5 * torch.sin(0.2 * x) * torch.sin(20 * x)\n```\n\nNow, instead of estimating the mean fitness, we can calculate a statistic that reflects the diversity of sampled behaviors. The variance of a distribution is one metric of diversity, and one variant of evolvability ES measures and optimizes such variance of behaviors sampled from the population distribution:\n\n```python\nsample = p.sample(n=100)\nbehaviors = behavior(sample)\nzscore = (behaviors - behaviors.mean()) / behaviors.std()\nvariance = expectation(zscore ** 2, sample, p=p)\n```\n\n### Maximizing Entropy\nIn the previous example, the gradient would be relatively straightforward to compute by hand.\nHowever, sometimes we may need to maximize objectives whose derivatives would be much more challenging to derive.\nIn particular, this final example will seek to maximize the entropy of the distribution of behaviors (a variant of evolvability ES).\n\nNote that for this example you'll also have to install `scipy` from pip.\n\nTo create a differentiable estimate of entropy, we first compute the pairwise distances between the different behaviors.\nNext, we create a smooth probability distribution by fitting a [kernel density estimate](https://en.wikipedia.org/wiki/Kernel_density_estimation):\n\n```python\ndists = scipy.spatial.distance.squareform(scipy.spatial.distance.pdist(behaviors, \"sqeuclidean\"))\nkernel = torch.tensor(scipy.exp(-dists / k_sigma ** 2), dtype=torch.float32)\np_x = expectation(kernel, sample, p=p, dim=1)\n```\n\nThen, we can use these probabilities to estimate the [entropy of the distribution](https://en.wikipedia.org/wiki/Entropy_(information_theory)), and run gradient descent on it as before:\n\n```python\nentropy = expectation(-torch.log(p_x), sample, p=p)\n```\n\nFull code for these examples can be found in the `demos` directory of\nthis repository.\n\n## Installation\nEither install EvoGrad from pip:\n```\npip install evograd\n```\nOr from the source code in this repository:\n```\ngit clone github.com/uber-research/EvoGrad\ncd EvoGrad\npip install -r requirements.txt\npip install -e .\n```\n\n\n## About\n\nDevelopment of EvoGrad was led by [Alex Gajewski](https://github.com/agajews) as a Summer intern at Uber AI Labs.\n"
  },
  {
    "path": "demos/cartpole.py",
    "content": "# Copyright (c) 2019 Uber Technologies, Inc.\n#\n# Licensed under the Uber Non-Commercial License (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at the root directory of this project.\n#\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\nimport torch\nimport numpy as np\n\nfrom evograd import expectation\nfrom evograd.distributions import Normal\n\nimport gym\n\n\ndef simulate_single(weights):\n    total_reward = 0.0\n    num_run = 10\n    for t in range(num_run):\n        observation = env.reset()\n        for i in range(300):\n            action = 1 if np.dot(weights, observation) > 0 else 0\n            observation, reward, done, info = env.step(action)\n            total_reward += reward\n            if done:\n                break\n    return total_reward / num_run\n\n\ndef simulate(batch_weights):\n    rewards = []\n    for weights in batch_weights:\n        rewards.append(simulate_single(weights.numpy()))\n    return torch.tensor(rewards)\n\n\nmu = torch.randn(4, requires_grad=True)  # population mean\nnpop = 50  # population size\nstd = 0.5  # noise standard deviation\nalpha = 0.03  # learning rate\np = Normal(mu, std)\nenv = gym.make(\"CartPole-v0\")\n\nfor t in range(2000):\n    sample = p.sample(npop)\n    fitnesses = simulate(sample)\n    scaled_fitnesses = (fitnesses - fitnesses.mean()) / fitnesses.std()\n    mean = expectation(scaled_fitnesses, sample, p=p)\n    mean.backward()\n\n    with torch.no_grad():\n        mu += alpha * mu.grad\n        mu.grad.zero_()\n\n    print(\"step: {}, mean fitness: {:0.5}\".format(t, float(fitnesses.mean())))\n"
  },
  {
    "path": "demos/max_ent_interference.py",
    "content": "# Copyright (c) 2019 Uber Technologies, Inc.\n#\n# Licensed under the Uber Non-Commercial License (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at the root directory of this project.\n#\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\nimport scipy\nimport scipy.spatial\nimport torch\nimport numpy as np\n\nfrom evograd import expectation\nfrom evograd.distributions import Normal\n\n\ndef fun(x):\n    return 5 * torch.sin(0.2 * x) * torch.sin(20 * x)\n\n\nmu = torch.tensor([1.0], requires_grad=True)\nnpop = 500  # population size\nstd = 0.5  # noise standard deviation\nk_sigma = 1.0  # kernel standard deviation\nalpha = 0.10  # learning rate\np = Normal(mu, std)\n\nfor t in range(2000):\n    sample = p.sample(npop)\n    novelties = fun(sample)\n    novelties = (novelties - novelties.mean()) / novelties.std()\n    dists = scipy.spatial.distance.squareform(\n        scipy.spatial.distance.pdist(novelties, \"sqeuclidean\")\n    )\n    kernel = torch.tensor(scipy.exp(-dists / k_sigma ** 2), dtype=torch.float32)\n    p_x = expectation(kernel, sample, p=p)\n    entropy = expectation(-torch.log(p_x), sample, p=p)\n    entropy.backward()\n\n    with torch.no_grad():\n        mu += alpha * mu.grad\n        mu.grad.zero_()\n\n    print(\"step: {}, estimated entropy: {:0.5}\".format(t, float(mu)))\n"
  },
  {
    "path": "demos/max_var_interference.py",
    "content": "# Copyright (c) 2019 Uber Technologies, Inc.\n#\n# Licensed under the Uber Non-Commercial License (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at the root directory of this project.\n#\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\nimport torch\nimport numpy as np\n\nfrom evograd import expectation\nfrom evograd.distributions import Normal\n\n\ndef fun(x):\n    return 5 * torch.sin(0.2 * x) * torch.sin(20 * x)\n\n\nmu = torch.tensor([1.0], requires_grad=True)\nnpop = 500  # population size\nstd = 0.5  # noise standard deviation\nalpha = 0.03  # learning rate\np = Normal(mu, std)\n\nfor t in range(2000):\n    sample = p.sample(npop)\n    behaviors = fun(sample)\n    zscores = (behaviors - behaviors.mean()) / behaviors.std()\n    variance = expectation(zscores ** 2, sample, p=p)\n    variance.backward()\n\n    with torch.no_grad():\n        mu += alpha * mu.grad\n        mu.grad.zero_()\n\n    print(\"step: {}, estimated variance: {:0.5}\".format(t, float(mu)))\n"
  },
  {
    "path": "demos/standard_interference.py",
    "content": "#Copyright (c) 2019 Uber Technologies, Inc.\n#\n#Licensed under the Uber Non-Commercial License (the \"License\");\n#you may not use this file except in compliance with the License.\n#You may obtain a copy of the License at the root directory of this project. \n#\n#See the License for the specific language governing permissions and\n#limitations under the License.\n\n\nimport torch\n\nfrom evograd import expectation\nfrom evograd.distributions import Normal\n\n\ndef fun(x):\n    return 5 * torch.sin(0.2 * x) * torch.sin(20 * x)\n\n\nmu = torch.tensor([1.0], requires_grad=True)\nnpop = 500  # population size\nstd = 0.5  # noise standard deviation\nalpha = 0.03  # learning rate\np = Normal(mu, std)\n\nfor t in range(2000):\n    sample = p.sample(npop)\n    fitnesses = fun(sample)\n    fitnesses = (fitnesses - fitnesses.mean()) / fitnesses.std()\n    mean = expectation(fitnesses, sample, p=p)\n    mean.backward()\n\n    with torch.no_grad():\n        mu += alpha * mu.grad\n        mu.grad.zero_()\n\n    print('step: {}, mean fitness: {:0.5}'.format(t, float(mu)))\n"
  },
  {
    "path": "demos/standard_quadratic.py",
    "content": "#Copyright (c) 2019 Uber Technologies, Inc.\n#\n#Licensed under the Uber Non-Commercial License (the \"License\");\n#you may not use this file except in compliance with the License.\n#You may obtain a copy of the License at the root directory of this project. \n#\n#See the License for the specific language governing permissions and\n#limitations under the License.\n\n\nimport torch\n\nfrom evograd import expectation\nfrom evograd.distributions import Normal\n\n\ndef fun(x):\n    return -(x - 5.0) ** 2\n\n\nmu = torch.tensor([1.0], requires_grad=True)\nnpop = 500  # population size\nstd = 0.5  # noise standard deviation\nalpha = 0.03  # learning rate\np = Normal(mu, std)\n\nfor t in range(2000):\n    sample = p.sample(npop)\n    fitnesses = fun(sample)\n    fitnesses = (fitnesses - fitnesses.mean()) / fitnesses.std()\n    mean = expectation(fitnesses, sample, p=p)\n    mean.backward()\n\n    with torch.no_grad():\n        mu += alpha * mu.grad\n        mu.grad.zero_()\n\n    print('step: {}, mean fitness: {:0.5}'.format(t, float(mu)))\n"
  },
  {
    "path": "evograd/__init__.py",
    "content": "#Copyright (c) 2019 Uber Technologies, Inc.\n#\n#Licensed under the Uber Non-Commercial License (the \"License\");\n#you may not use this file except in compliance with the License.\n#You may obtain a copy of the License at the root directory of this project. \n#\n#See the License for the specific language governing permissions and\n#limitations under the License.\n\n\nfrom .expectation import expectation\n"
  },
  {
    "path": "evograd/distributions.py",
    "content": "# Copyright (c) 2019 Uber Technologies, Inc.\n#\n# Licensed under the Uber Non-Commercial License (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at the root directory of this project.\n#\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\nimport numpy as np\nimport torch\n\nfrom .noise import noise\n\n\nclass NormalProbRatio(torch.autograd.Function):\n    @staticmethod\n    def forward(ctx, mu, sigma, descriptors, decode_fn):\n        ctx.save_for_backward(mu)\n        ctx.sigma = sigma\n        ctx.descriptors = descriptors\n        ctx.decode_fn = decode_fn\n        res = torch.ones(len(descriptors), dtype=torch.float32)\n        res.requires_grad = True\n        return res\n\n    @staticmethod\n    def backward(ctx, grad_output):\n        mu, = ctx.saved_tensors\n        theta = ctx.decode_fn(ctx.descriptors)\n        grad = (theta - mu) / ctx.sigma ** 2 * grad_output.unsqueeze(1)\n        return (grad, None, None, None)\n\n\nclass MixNormalProbRatio(torch.autograd.Function):\n    @staticmethod\n    def forward(ctx, mus, sigma, descriptors, decode_fn):\n        ctx.save_for_backward(mus)\n        ctx.sigma = sigma\n        ctx.descriptors = descriptors\n        ctx.decode_fn = decode_fn\n        res = torch.ones(len(descriptors), dtype=torch.float32)\n        res.requires_grad = True\n        return res\n\n    @staticmethod\n    def backward(ctx, grad_output):\n        mus, = ctx.saved_tensors\n        thetas = ctx.decode_fn(ctx.descriptors)\n        epsilons = [thetas - mu for mu in mus]\n        grads = torch.stack(\n            [\n                (epsilon / ctx.sigma ** 2)\n                / (\n                    1\n                    + sum(\n                        torch.exp(\n                            -0.5\n                            * (other.dot(other) - epsilon.dot(epsilon))\n                            / ctx.sigma ** 2\n                        )\n                        for other in epsilons\n                        if other is not epsilon\n                    )\n                )\n                * grad_output\n                for epsilon in epsilons\n            ]\n        )\n        return (grads, None, None, None)\n\n\nclass Distribution:\n    def __init__(self, device, random_state=None):\n        if random_state is None:\n            random_state = np.random.RandomState()  # pylint: disable=no-member\n        self.random_state = random_state\n        self.device = device\n\n\nclass Normal(Distribution):\n    def __init__(self, mu, sigma, random_state=None):\n        \"\"\"\n        mu: torch.tensor\n        sigma: torch.tensor or float\n        \"\"\"\n        super().__init__(mu.device, random_state)\n        self.mu = mu\n        self.sigma = sigma\n\n    def ratio(self, descriptors):\n        return NormalProbRatio.apply(self.mu, self.sigma, descriptors, self.decode)\n\n    def sample(self, n, encode=False):\n        n_epsilons = n\n        noise_inds = np.asarray(\n            [\n                noise.sample_index(self.random_state, len(self.mu))\n                for _ in range(n_epsilons)\n            ],\n            dtype=\"int\",\n        )\n        descriptors = [(idx, 1) for idx in noise_inds]\n        if encode:\n            return descriptors\n        thetas = torch.stack([self.decode(descriptor) for descriptor in descriptors])\n        return thetas\n\n    def decode(self, descriptor):\n        if not isinstance(descriptor, tuple):\n            # assert isinstance(descriptor, torch.tensor)\n            return descriptor\n        noise_idx, direction = descriptor\n        epsilon = torch.tensor(noise.get(noise_idx, len(self.mu)), device=self.device)\n        with torch.no_grad():\n            return self.mu + direction * self.sigma * epsilon\n\n\nclass PairedNormal(Normal):\n    def sample(self, n, encode=False):\n        assert n % 2 == 0\n        n_epsilons = n // 2\n        noise_inds = np.asarray(\n            [\n                noise.sample_index(self.random_state, len(self.mu))\n                for _ in range(n_epsilons)\n            ],\n            dtype=\"int\",\n        )\n        descriptors = [(idx, 1) for idx in noise_inds] + [\n            (idx, -1) for idx in noise_inds\n        ]\n        if encode:\n            return descriptors\n        thetas = torch.stack([self.decode(descriptor) for descriptor in descriptors])\n        return thetas\n\n\nclass MixNormal(Distribution):\n    def __init__(self, mus, sigma, random_state=None):\n        \"\"\"\n        mu: torch.tensor\n        sigma: torch.tensor or float\n        \"\"\"\n        super().__init__(mus[0].device, random_state)\n        self.mus = [mu for mu in mus]\n        self.sigma = sigma\n\n    def ratio(self, descriptor):\n        return MixNormalProbRatio.apply(self.mus, self.sigma, descriptor, self.decode)\n\n    def sample(self, n, encode=False):\n        n_epsilons = n\n        noise_ids = np.asarray(\n            [\n                noise.sample_index(self.random_state, len(self.mus[0]))\n                for _ in range(n_epsilons)\n            ],\n            dtype=\"int\",\n        )\n        mu_ids = np.random.randint(len(self.mus), size=n)\n        descriptors = [\n            (noise_id, mu_id, 1) for noise_id, mu_id in zip(noise_ids, mu_ids)\n        ]\n        if encode:\n            return descriptors\n        thetas = torch.stack([self.decode(descriptor) for descriptor in descriptors])\n        return thetas\n\n    def decode(self, descriptor):\n        if not isinstance(descriptor, tuple):\n            # assert isinstance(descriptor, torch.tensor)\n            return descriptor\n        noise_id, mu_id, direction = descriptor\n        epsilon = torch.tensor(\n            noise.get(noise_id, len(self.mus[0])), device=self.device\n        )\n        with torch.no_grad():\n            return self.mus[mu_id] + direction * self.sigma * epsilon\n"
  },
  {
    "path": "evograd/expectation.py",
    "content": "# Copyright (c) 2019 Uber Technologies, Inc.\n#\n# Licensed under the Uber Non-Commercial License (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at the root directory of this project.\n#\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\nimport torch\n\n\ndef unsqueeze_as(x, y):\n    assert len(y.shape) >= len(x.shape)\n    for _ in range(len(y.shape) - len(x.shape)):\n        x = torch.unsqueeze(x, -1)\n    return x\n\n\ndef expectation(vals, sample, p):\n    ratio = unsqueeze_as(p.ratio(sample), vals)\n    prod = vals * ratio\n    return prod.mean(dim=0)\n"
  },
  {
    "path": "evograd/noise.py",
    "content": "#Copyright (c) 2019 Uber Technologies, Inc.\n#\n#Licensed under the Uber Non-Commercial License (the \"License\");\n#you may not use this file except in compliance with the License.\n#You may obtain a copy of the License at the root directory of this project. \n#\n#See the License for the specific language governing permissions and\n#limitations under the License.\n\n\nimport logging\n\nimport numpy as np\n\nlogger = logging.getLogger(__name__)\nlogging.basicConfig(level=logging.INFO)\n\ndebug = True\n\n\nclass SharedNoiseTable:\n    def __init__(self):\n        import ctypes\n        import multiprocessing\n\n        seed = 42\n        # 1 gigabyte of 32-bit numbers. Will actually sample 2 gigabytes below.\n        count = 250000000 if not debug else 10000000\n        logger.info(\"Sampling {} random numbers with seed {}\".format(count, seed))\n        self._shared_mem = multiprocessing.Array(ctypes.c_float, count)\n        self.noise = np.ctypeslib.as_array(self._shared_mem.get_obj())\n        assert self.noise.dtype == np.float32\n        self.noise[:] = np.random.RandomState(seed).randn(  # pylint: disable=no-member\n            count\n        )  # 64-bit to 32-bit conversion here\n        logger.info(\"Sampled {} bytes\".format(self.noise.size * 4))\n\n    def get(self, i, dim):\n        return self.noise[i : i + dim]\n\n    def sample_index(self, stream, dim):\n        return stream.randint(0, len(self.noise) - dim + 1)\n\n\nnoise = SharedNoiseTable()\n"
  },
  {
    "path": "requirements.txt",
    "content": "numpy==1.16.4\ntorch==1.1.0.post2\n"
  },
  {
    "path": "setup.py",
    "content": "# Copyright (c) 2019 Uber Technologies, Inc.\n#\n# Licensed under the Uber Non-Commercial License (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at the root directory of this project.\n#\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\nimport setuptools\n\nwith open(\"README.md\", \"r\") as fh:\n    long_description = fh.read()\n\nsetuptools.setup(\n    name=\"evograd\",\n    version=\"0.1.2\",\n    author=\"Alex Gajewski\",\n    author_email=\"agajews@gmail.com\",\n    description=\"A lightweight tool for differentiating through expectations\",\n    long_description=long_description,\n    long_description_content_type=\"text/markdown\",\n    url=\"https://github.com/uber-research/EvoGrad\",\n    packages=setuptools.find_packages(),\n    install_requires=[\"numpy\", \"torch\"],\n    classifiers=[\n        \"Programming Language :: Python :: 3\",\n        \"License :: OSI Approved :: MIT License\",\n        \"Operating System :: OS Independent\",\n    ],\n)\n"
  }
]