bce31067afb1fa3fd37db92ad2b82db0440af183,thinc/backends/jax_ops.py,JaxOps,max_pool,#JaxOps#Any#Any#,165

Before Change

        return Y

    def max_pool(self, X: Array2d, lengths: Array1d) -> Array2d:
        Y = self.alloc_f2d(lengths.shape[0], X.shape[1])
        start = 0
        for i, length in enumerate(lengths):
            Y[i] = X[start : start + length].max(axis=0)
            start += length
        return Y

    def backprop_sum_pool(self, d_sums: Array2d, lengths: Array1d) -> Array2d:
        dX = self.alloc_f2d(lengths.sum(), d_sums.shape[1])
        start = 0

After Change

        return mean_pool(X, lengths)

    def max_pool(self, X: Array2d, lengths: Array1d) -> Array2d:
        return max_pool(X, lengths)

    def backprop_sum_pool(self, d_sums: Array2d, lengths: Array1d) -> Array2d:
        return backprop_sum_pool(d_sums, lengths)

Instances