fix: Fix problem using ops.ones in BigBird (#1861)

reeered · web-flow · commit 1a66f9f82e9b · 2024-12-11T15:41:51.000+08:00
diff --git a/mindnlp/transformers/models/big_bird/modeling_big_bird.py b/mindnlp/transformers/models/big_bird/modeling_big_bird.py
@@ -470,13 +470,13 @@ def bigbird_block_sparse_attention(
             [
                 to_mask[:, :, :, : 3 * to_block_size],
                 to_mask[:, :, :, -to_block_size:],
-                ops.ones([bsz, 1, 1, n_rand_blocks * to_block_size], dtype=to_mask.dtype),
+                ops.ones((bsz, 1, 1, n_rand_blocks * to_block_size), dtype=to_mask.dtype),
             ],
             dim=3,
         )
         second_rand_pad = ops.cat(
             [
-                ops.ones([bsz, n_heads, from_block_size, 4 * to_block_size], dtype=rand_mask.dtype),
+                ops.ones((bsz, n_heads, from_block_size, 4 * to_block_size), dtype=rand_mask.dtype),
                 rand_mask[:, :, 0],
             ],
             dim=3,
@@ -604,13 +604,13 @@ def bigbird_block_sparse_attention(
             [
                 to_mask[:, :, :, :to_block_size],
                 to_mask[:, :, :, -3 * to_block_size :],
-                ops.ones([bsz, 1, 1, n_rand_blocks * to_block_size], dtype=to_mask.dtype),
+                ops.ones((bsz, 1, 1, n_rand_blocks * to_block_size), dtype=to_mask.dtype),
             ],
             dim=3,
         )
         second_last_rand_pad = ops.cat(
             [
-                ops.ones([bsz, n_heads, from_block_size, 4 * to_block_size], dtype=rand_mask.dtype),
+                ops.ones((bsz, n_heads, from_block_size, 4 * to_block_size), dtype=rand_mask.dtype),
                 rand_mask[:, :, -1],
             ],
             dim=3,
diff --git a/mindnlp/transformers/models/bigbird_pegasus/modeling_bigbird_pegasus.py b/mindnlp/transformers/models/bigbird_pegasus/modeling_bigbird_pegasus.py
@@ -450,13 +450,13 @@ def bigbird_block_sparse_attention(
             [
                 to_mask[:, :, :, : 3 * to_block_size],
                 to_mask[:, :, :, -to_block_size:],
-                ops.ones([bsz, 1, 1, n_rand_blocks * to_block_size], dtype=to_mask.dtype),
+                ops.ones((bsz, 1, 1, n_rand_blocks * to_block_size), dtype=to_mask.dtype),
             ],
             dim=3,
         )
         second_rand_pad = ops.cat(
             [
-                ops.ones([bsz, n_heads, from_block_size, 4 * to_block_size], dtype=rand_mask.dtype),
+                ops.ones((bsz, n_heads, from_block_size, 4 * to_block_size), dtype=rand_mask.dtype),
                 rand_mask[:, :, 0],
             ],
             dim=3,
@@ -584,13 +584,13 @@ def bigbird_block_sparse_attention(
             [
                 to_mask[:, :, :, :to_block_size],
                 to_mask[:, :, :, -3 * to_block_size :],
-                ops.ones([bsz, 1, 1, n_rand_blocks * to_block_size], dtype=to_mask.dtype),
+                ops.ones((bsz, 1, 1, n_rand_blocks * to_block_size), dtype=to_mask.dtype),
             ],
             dim=3,
         )
         second_last_rand_pad = ops.cat(
             [
-                ops.ones([bsz, n_heads, from_block_size, 4 * to_block_size], dtype=rand_mask.dtype),
+                ops.ones((bsz, n_heads, from_block_size, 4 * to_block_size), dtype=rand_mask.dtype),
                 rand_mask[:, :, -1],
             ],
             dim=3,

Original file line number	Diff line number	Diff line change
`@@ -470,13 +470,13 @@ def bigbird_block_sparse_attention(`
`470`	`470`	`[`
`471`	`471`	`to_mask[:, :, :, : 3 * to_block_size],`
`472`	`472`	`to_mask[:, :, :, -to_block_size:],`
`473`		`- ops.ones([bsz, 1, 1, n_rand_blocks * to_block_size], dtype=to_mask.dtype),`
	`473`	`+ ops.ones((bsz, 1, 1, n_rand_blocks * to_block_size), dtype=to_mask.dtype),`
`474`	`474`	`],`
`475`	`475`	`dim=3,`
`476`	`476`	`)`
`477`	`477`	`second_rand_pad = ops.cat(`
`478`	`478`	`[`
`479`		`- ops.ones([bsz, n_heads, from_block_size, 4 * to_block_size], dtype=rand_mask.dtype),`
	`479`	`+ ops.ones((bsz, n_heads, from_block_size, 4 * to_block_size), dtype=rand_mask.dtype),`
`480`	`480`	`rand_mask[:, :, 0],`
`481`	`481`	`],`
`482`	`482`	`dim=3,`
`@@ -604,13 +604,13 @@ def bigbird_block_sparse_attention(`
`604`	`604`	`[`
`605`	`605`	`to_mask[:, :, :, :to_block_size],`
`606`	`606`	`to_mask[:, :, :, -3 * to_block_size :],`
`607`		`- ops.ones([bsz, 1, 1, n_rand_blocks * to_block_size], dtype=to_mask.dtype),`
	`607`	`+ ops.ones((bsz, 1, 1, n_rand_blocks * to_block_size), dtype=to_mask.dtype),`
`608`	`608`	`],`
`609`	`609`	`dim=3,`
`610`	`610`	`)`
`611`	`611`	`second_last_rand_pad = ops.cat(`
`612`	`612`	`[`
`613`		`- ops.ones([bsz, n_heads, from_block_size, 4 * to_block_size], dtype=rand_mask.dtype),`
	`613`	`+ ops.ones((bsz, n_heads, from_block_size, 4 * to_block_size), dtype=rand_mask.dtype),`
`614`	`614`	`rand_mask[:, :, -1],`
`615`	`615`	`],`
`616`	`616`	`dim=3,`
Original file line number	Diff line number	Diff line change
`@@ -450,13 +450,13 @@ def bigbird_block_sparse_attention(`
`450`	`450`	`[`
`451`	`451`	`to_mask[:, :, :, : 3 * to_block_size],`
`452`	`452`	`to_mask[:, :, :, -to_block_size:],`
`453`		`- ops.ones([bsz, 1, 1, n_rand_blocks * to_block_size], dtype=to_mask.dtype),`
	`453`	`+ ops.ones((bsz, 1, 1, n_rand_blocks * to_block_size), dtype=to_mask.dtype),`
`454`	`454`	`],`
`455`	`455`	`dim=3,`
`456`	`456`	`)`
`457`	`457`	`second_rand_pad = ops.cat(`
`458`	`458`	`[`
`459`		`- ops.ones([bsz, n_heads, from_block_size, 4 * to_block_size], dtype=rand_mask.dtype),`
	`459`	`+ ops.ones((bsz, n_heads, from_block_size, 4 * to_block_size), dtype=rand_mask.dtype),`
`460`	`460`	`rand_mask[:, :, 0],`
`461`	`461`	`],`
`462`	`462`	`dim=3,`
`@@ -584,13 +584,13 @@ def bigbird_block_sparse_attention(`
`584`	`584`	`[`
`585`	`585`	`to_mask[:, :, :, :to_block_size],`
`586`	`586`	`to_mask[:, :, :, -3 * to_block_size :],`
`587`		`- ops.ones([bsz, 1, 1, n_rand_blocks * to_block_size], dtype=to_mask.dtype),`
	`587`	`+ ops.ones((bsz, 1, 1, n_rand_blocks * to_block_size), dtype=to_mask.dtype),`
`588`	`588`	`],`
`589`	`589`	`dim=3,`
`590`	`590`	`)`
`591`	`591`	`second_last_rand_pad = ops.cat(`
`592`	`592`	`[`
`593`		`- ops.ones([bsz, n_heads, from_block_size, 4 * to_block_size], dtype=rand_mask.dtype),`
	`593`	`+ ops.ones((bsz, n_heads, from_block_size, 4 * to_block_size), dtype=rand_mask.dtype),`
`594`	`594`	`rand_mask[:, :, -1],`
`595`	`595`	`],`
`596`	`596`	`dim=3,`