title: Sequence of for loops to access different dimensions of a tensor
body: <p>My problem is a bit hard to explain.</p>
  <ul>
  <li>I want to update probability distributions recursively in [0, 1]^n.</li>
  <li>I discretized it with 10 bins per dimension, so there are 10^n cells indexed by an n-tuple.</li>
  <li>Initially, I should have a uniform distribution; but then I take a threshold parameter that dictates how I propagate the distribution: every point that is &quot;below&quot; all the n thresholds must receive zero probability.</li>
  </ul>
  <p>Here is how I hardcoded it for 1 or 3 dimensions:</p>
  <pre><code>import torch
  import numpy as np
  
  dim = 3 # or 1
  bins = 10
  
  
  # form the vector [10, 10, ..., 10] (n dimensions)
  probabilityStarter = []
  for counter in range(dim):
      probabilityStarter.append(bins)
  
  
  # form the n-cube with 10^n cells
  probabilityDistribution = torch.ones(probabilityStarter, dtype = torch.int8)
  zeroProbability = torch.ones(probabilityStarter, dtype = torch.int8)
  
  
  # if lower than the threshold, I want to zero the probability
  threshold = 0.5
  
  
  # for 1 dim, I just need to re-arrange, skewing the probabilities &quot;to the right&quot;
  if dim == 1:
      for j in range(bins):
          step = (j + 0.5) / bins
          if step &lt; threshold:
              probabilityDistribution[j] = 0
      
      probabilityDistribution = probabilityDistribution.float()
      probabilityDistribution = probabilityDistribution / torch.sum(probabilityDistribution)
  
  
  # for n dim, I need to do n for loops and update the d-th dimension [:, ..., d, :, ..., :]
  elif dim == 3:
      for j in range(bins):
          step = (j + 0.5) / bins
          if step &gt; threshold:
              zeroProbability[j, :, :] = 0 # for dim 1
  
      for j in range(bins):
          step = (j + 0.5) / bins
          if step &gt; threshold:
              zeroProbability[:, j, :] = 0 # for dim 2
  
      for j in range(bins):
          step = (j + 0.5) / bins
          if step &gt; threshold:
              zeroProbability[:, :, j] = 0 # for dim 3
  
      probabilityDistribution = probabilityDistribution - zeroProbability
      probabilityDistribution = probabilityDistribution.float()
      probabilityDistribution = probabilityDistribution / torch.sum(probabilityDistribution)
  </code></pre>
  <p>However, if I want the 'dim' to be dynamic, how could I implement it?</p>