Performs a forward pass using multiple GPUs.  This is a simplification
    of torch.nn.parallel.data_parallel to support the allennlp model
    interface.
737
def data_parallel(batch_group: List[TensorDict],
                  model: Model,
                  cuda_devices: List) -> Dict[str, torch.Tensor]:
    
    assert len(batch_group) <= len(cuda_devices)

    moved = [nn_util.move_to_device(batch, device)
             for batch, device in zip(batch_group, cuda_devices)]

    used_device_ids = cuda_devices[:len(moved)]
    # Counterintuitively, it appears replicate expects the source device id to be the first element
    # in the device id list. See torch.cuda.comm.broadcast_coalesced, which is called indirectly.
    replicas = replicate(model, used_device_ids)

    # We pass all our arguments as kwargs. Create a list of empty tuples of the
    # correct shape to serve as (non-existent) positional arguments.
    inputs = [()] * len(batch_group)
    outputs = parallel_apply(replicas, inputs, moved, used_device_ids)

    # Only the 'loss' is needed.
    # a (num_gpu, ) tensor with loss on each GPU
    losses = gather([output['loss'].unsqueeze(0) for output in outputs], used_device_ids[0], 0)
    return {'loss': losses.mean()}
266323
def to_gpu(*args):
    
    if len(args) > 1:
        return (cp.asarray(x) for x in args)
    else:
        return cp.asarray(args[0])
348562
def _learner_distributed(learn:Learner, cuda_id:int, cache_dir:PathOrStr='tmp'):
    ""
    learn.callbacks.append(DistributedTrainer(learn, cuda_id))
    learn.callbacks.append(DistributedRecorder(learn, cuda_id, cache_dir))
    return learn
1100507
def forward(self, actions, batch_info):
        
        while len(self.processes) < actions.shape[0]:
            len_action_space = self.action_space.shape[-1]

            self.processes.append(
                OrnsteinUhlenbeckNoiseProcess(
                    np.zeros(len_action_space), float(self.std_dev) * np.ones(len_action_space)
                )
            )

        noise = torch.from_numpy(np.stack([x() for x in self.processes])).float().to(actions.device)

        return torch.min(torch.max(actions + noise, self.low_tensor), self.high_tensor)
23642
def main(
    gpus:Param("The GPUs to use for distributed training", str)='all',
    script:Param("Script to run", str, opt=False)='',
    args:Param("Args to pass to script", nargs='...', opt=False)=''
):
    ""
    # Loosely based on torch.distributed.launch
    current_env = os.environ.copy()
    gpus = list(range(torch.cuda.device_count())) if gpus=='all' else list(gpus)
    current_env["WORLD_SIZE"] = str(len(gpus))
    current_env["MASTER_ADDR"] = '127.0.0.1'
    current_env["MASTER_PORT"] = '29500'

    processes = []
    for i,gpu in enumerate(gpus):
        current_env["RANK"] = str(i)
        cmd = [sys.executable, "-u", script, f"--gpu={gpu}"] + args
        process = subprocess.Popen(cmd, env=current_env)
        processes.append(process)

    for process in processes: process.wait()
637345
def call_for_each_tower(self, tower_fn):
        
        # if tower_fn returns [(grad, var), ...], this returns #GPU x #VAR x 2
        return DataParallelBuilder.build_on_towers(
            self.towers,
            tower_fn,
            # use no variable scope for the first tower
            use_vs=[False] + [True] * (len(self.towers) - 1))
772109
def run_forward_backward(self, batch: mx.io.DataBatch, metric: mx.metric.EvalMetric):
        
        self.module.forward_backward(batch)
        self.module.update_metric(metric, batch.label)
675434
def forward_backward(self, parameters):
        
        # If the features are sparse, we can use an optimization.
        # I'm not using scipy.sparse here because we want to avoid a scipy dependency and also scipy.sparse doesn't seem
        # to handle arrays of shape higher than 2.
        if isinstance(self.sparse_x, str) and self.sparse_x == 'uninitialized':
            if (self.x == 0).sum() * 1.0 / self.x.size > 0.6:
                self.sparse_x = self._construct_sparse_features(self.x)
            else:
                self.sparse_x = 'not sparse'

        I, J, K = self.x.shape
        if not isinstance(self.sparse_x, str):
            C = self.sparse_x[0].shape[2]
            S, _ = parameters.shape
            x_dot_parameters = np.zeros((I, J, S))
            sparse_multiply(x_dot_parameters, self.sparse_x[0], self.sparse_x[1], parameters.T, I, J, K, C, S)
        else:
            x_dot_parameters = np.dot(self.x, parameters.T)  # Pre-compute the dot product
        alpha = self._forward(x_dot_parameters)
        beta = self._backward(x_dot_parameters)
        classes_to_ints = {k: i for i, k in enumerate(set(self.states_to_classes.values()))}
        states_to_classes = np.array([classes_to_ints[self.states_to_classes[state]]
                                      for state in range(max(self.states_to_classes.keys()) + 1)], dtype='int64')
        if not isinstance(self.sparse_x, str):
            ll, deriv = gradient_sparse(alpha, beta, parameters, states_to_classes,
                                        self.sparse_x[0], self.sparse_x[1], classes_to_ints[self.y],
                                        I, J, self.sparse_x[0].shape[2])
        else:
            ll, deriv = gradient(alpha, beta, parameters, states_to_classes,
                                 self.x, classes_to_ints[self.y], I, J, K)
        return ll, deriv
99187
function parallel (fns, next) {
  var len = fns.length;
  var values = [];
  var skip = false;
  function done(i, e, v) {
    values[i] = v;
    len--;
    if (!skip && (e || !len)) {
      skip = true;
      next(e, e ? null : values);
    }
  }
  fns.forEach(function (fn, i) {
    fn(function(e, v){done(i, e, v);});
  });
}
19521
def sample_forward_transitions(self, batch_size, batch_info, forward_steps: int,
                                   discount_factor: float) -> Transitions:
        
        indexes = self.backend.sample_batch_transitions(batch_size, forward_steps=forward_steps)
        transition_tensors = self.backend.get_transitions_forward_steps(
            indexes, forward_steps=forward_steps, discount_factor=discount_factor
        )

        return Trajectories(
            num_steps=batch_size,
            num_envs=self.backend.num_envs,
            environment_information=None,
            transition_tensors={k: torch.from_numpy(v) for k, v in transition_tensors.items()},
            rollout_tensors={},
            extra_data={
                'forward_steps': forward_steps
            }
        ).to_transitions()
1170155
def _forwardImplementation(self, inbuf, outbuf):
        
        assert self.module

        propensities = self.module.getActionValues(0)

        summedProps = sum(propensities)
        probabilities = propensities / summedProps

        action = eventGenerator(probabilities)
#        action = drawIndex(probabilities)

        outbuf[:] = scipy.array([action])
1036297
def forward(self, data_batch, is_train=None):
        
        assert self.binded and self.params_initialized

        # make a shallow copy, just to maintain necessary properties (if any) like
        # bucket_key, pad, etc.
        data_batch = copy.copy(data_batch)

        for i_layer, module in enumerate(self._modules):
            module.forward(data_batch, is_train=is_train)

            if i_layer+1 == len(self._modules):
                # the last layer, do not need to do the followings
                break

            data_batch.data = module.get_outputs()
            if hasattr(data_batch, 'provide_data'):
                # need to update this, in case the internal module is using bucketing
                # or whatever
                data_names = [x[0] for x in module.output_shapes]
                assert len(data_names) == len(data_batch.data)
                data_batch.provide_data = [(name, x.shape) for name, x in
                                           zip(data_names, data_batch.data)]
931642
def begin_update(self, x_data, drop=0.0):
        
        x_var = torch.autograd.Variable(xp2torch(x_data), requires_grad=True)
        # Make prediction

        y_var = self._model(x_var)

        def backward_pytorch(dy_data, sgd=None):
            dy_var = xp2torch(dy_data)
            torch.autograd.backward((y_var,), grad_tensors=(dy_var,))
            if sgd is not None:
                if self._optimizer is None:
                    self._optimizer = self._create_optimizer(sgd)
                self._optimizer.step()
                self._optimizer.zero_grad()
            return torch2xp(x_var.grad)

        return torch2xp(y_var), backward_pytorch
106848
@Override
    public Tensor forward() {
        Tensor x = modInX.getOutput();
        Tensor w = modInW.getOutput();
        y = new Tensor(x); // copy
        y.elemMultiply(w);
        return y;
    }
274472
def forward(self, x):
        
        if not self.training or not self.p:
            return x
        x = x.clone()
        mask = x.new_empty(1, x.size(1), x.size(2), requires_grad=False).bernoulli_(1 - self.p)
        mask = mask.div_(1 - self.p)
        mask = mask.expand_as(x)
        return x * mask
29853
def run(samples, run_parallel):
    
    to_process = []
    extras = []
    for data in (xs[0] for xs in samples):
        hlacaller = tz.get_in(["config", "algorithm", "hlacaller"], data)
        if hlacaller:
            to_process.append(data)
        else:
            extras.append([data])
    processed = run_parallel("call_hla", ([x] for x in to_process))
    return extras + processed
687596
def forward_prediction(self, m):
        
        return self.infer(self.conf.m_dims, self.conf.s_dims, m)